mx.nd.linalg.extractdiag

Description

Extracts the diagonal entries of a square matrix. Input is a tensor A of dimension n >= 2.

If n=2, then A represents a single square matrix which diagonal elements get extracted as a 1-dimensional tensor.

If n>2, then A represents a batch of square matrices on the trailing two dimensions. The extracted diagonals are returned as an n-1-dimensional tensor.

Note

The operator supports float32 and float64 data types only.

Example:

Single matrix diagonal extraction
A = [[1.0, 2.0],
[3.0, 4.0]]

extractdiag(A) = [1.0, 4.0]

extractdiag(A, 1) = [2.0]

Batch matrix diagonal extraction
A = [[[1.0, 2.0],
[3.0, 4.0]],
[[5.0, 6.0],
[7.0, 8.0]]]

extractdiag(A) = [[1.0, 4.0],
[5.0, 8.0]]

Arguments

Argument

Description

A

NDArray-or-Symbol.

Tensor of square matrices

offset

int, optional, default=’0’.

Offset of the diagonal versus the main diagonal. 0 corresponds to the main diagonal, a negative/positive value to diagonals below/above the main diagonal.