mindspore.ops.diag_embed

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

Create a tensor whose diagonals of certain 2D planes (specified by dim1 and dim2) are filled by input, and all other positions are set to 0. The 2D planes formed by the last two dimensions of the returned tensor are chosen by default.

Parameters

  • input (Tensor) – Values to fill diagonal.

  • 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 for diagonal filling. Default -2 .

  • dim2 (int, optional) – The second dimension for diagonal filling. Default -1 .

Returns

A tensor with the same dtype as input, and with shape that has one dimension higher than the input.

Raises

ValueError – If the dimension of input is not 1D-6D.

Supported Platforms:

Ascend GPU CPU

Examples

>>> import mindspore
>>> input = mindspore.tensor([[1, 2, 3],
...                           [4, 5, 6],
...                           [7, 8, 9]])
>>> mindspore.ops.diag_embed(input)
Tensor(shape=[3, 3, 3], dtype=Int64, value=
[[[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.diag_embed(input, offset=1, dim1=0, dim2=1)
Tensor(shape=[4, 4, 3], dtype=Int64, value=
[[[0, 0, 0], [1, 4, 7], [0, 0, 0], [0, 0, 0]],
 [[0, 0, 0], [0, 0, 0], [2, 5, 8], [0, 0, 0]],
 [[0, 0, 0], [0, 0, 0], [0, 0, 0], [3, 6, 9]],
 [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]])