mindspore.ops.diagonal

View Source On Gitee
mindspore.ops.diagonal(input, offset=0, dim1=0, dim2=1)[source]

Returns diagonals of the input tensor along specified dimension.

If input is 2-D, returns a 1-D tensor containing the diagonal of input with the given offset.

If input has more than two dimensions, then the diagonals of specified 2-D sub-array determined by dim1 and dim2 is returned. The shape of returned tensor is the original shape with axis1 and axis2 removed and a new dimension inserted at the end corresponding to the diagonal.

Parameters

  • input (Tensor) – The input tensor with at least two dimensions.

  • offset (int, optional) –

    Diagonal offset. Default 0 .

    • When offset is a positive integer, shift the diagonal upward.

    • When offset is a negative integer, shift the diagonal downward.

  • dim1 (int, optional) – The first dimension specifying the 2D plane. Default 0 .

  • dim2 (int, optional) – The second dimension specifying the 2D plane. Default 1 .

Returns

Tensor

Supported Platforms:

Ascend GPU CPU

Examples

>>> import mindspore
>>> input = mindspore.tensor([[[1, 0, 0],
...                            [0, 2, 0],
...                            [0, 0, 3]],
...                           [[4, 0, 0],
...                            [0, 5, 0],
...                            [0, 0, 6]],
...                           [[7, 0, 0],
...                            [0, 8, 0],
...                            [0, 0, 9]]])
>>> mindspore.ops.diagonal(input)
Tensor(shape=[3, 3], dtype=Int64, value=
[[1, 0, 0],
 [0, 5, 0],
 [0, 0, 9]])
>>> mindspore.ops.diagonal(input, offset=1)
Tensor(shape=[3, 2], dtype=Int64, value=
[[0, 0],
 [2, 0],
 [0, 6]])
>>> mindspore.ops.diagonal(input, offset=0, dim1=2, dim2=1)
Tensor(shape=[3, 3], dtype=Int64, value=
[[1, 2, 3],
 [4, 5, 6],
 [7, 8, 9]])