diagonal_op

→TensorImplPtr

int diagonal_op(const int & a, int offset, int axis1, int axis2)

source edit

Python wrappers

Extract a single diagonal as a final-axis vector.

Returns the elements of a lying on the diagonal of the (axis1, axis2) plane offset by offset: positive offsets shift the diagonal toward the upper-right, negative offsets toward the lower-left. The two selected axes are removed and a new trailing axis of length L is appended.

Parameters

aTensorImplPtr

Input tensor of rank

\ge 2

offsetint

Diagonal offset (0 = main diagonal).

axis1int

First axis defining the plane.

axis2int

Second axis defining the plane. Must differ from axis1; the two are canonicalised to a1 < a2.

Returns

TensorImplPtr

Tensor whose shape is a.shape with axes axis1 and axis2 removed and a trailing dimension of length $L = \max(0, \min(M - r_0, N - c_0))$ appended, where $M, N$ are the sizes of the two selected axes and $r_0, c_0$ are the starting row / column for the given offset.

Notes

Backward scatters the diagonal gradient into a zero tensor of the original input shape, leaving off-diagonal positions zero.