diagonal

→Tensor

diagonal(A: Tensor, offset: int = 0, dim1: int = -2, dim2: int = -1)

source

Extract a (off-)diagonal from each matrix in a batched input.

Returns the elements of $A$ lying on the diagonal selected by offset from the matrix slice formed by (dim1, dim2). For a 2-D matrix and offset=0 this is

d_i \,=\, A_{i,\,i}.

Positive offset selects super-diagonals ( $A_{i, i + \text{offset}}$ ); negative offsets select sub-diagonals.

Parameters

ATensor

Input of shape (*, m, n) (or higher rank).

offset(int, keyword - only)= 0

Diagonal index relative to the main diagonal. Default 0.

dim1(int, keyword - only)= -2

First matrix dimension. Default -2.

dim2(int, keyword - only)= -1

Second matrix dimension. Default -1.

Returns

Tensor

Diagonal values with the two matrix axes replaced by a single axis of length $\min(m, n) - |\text{offset}|$ .

Notes

Shares the engine kernel with the top-level lucid.diagonal; the only difference is the keyword-only / matrix-aware defaults (dim1=-2, dim2=-1) that match the standard linalg convention.

Examples

>>> import lucid
>>> from lucid.linalg import diagonal
>>> A = lucid.tensor([[1.0, 2.0], [3.0, 4.0]])
>>> diagonal(A)
Tensor([1.0000, 4.0000])
>>> diagonal(A, offset=1)
Tensor([2.0000])