mindspore.ops.diagonal_scatter

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

dim1 and dim2 specify the two dimensions of input, the elements in these two dimensions will be treated as elements of a matrix, and src is embedded on the diagonal of the matrix.

Parameters

  • input (Tensor) – Input Tensor, whose dimension is larger than 1.

  • src (Tensor) – The source Tensor to embed.

  • offset (int, optional) –

    offset controls which diagonal to choose. Default: 0 .

    • When offset is zero, the diagonal chosen is the main diagonal.

    • When offset is a positive integer, the diagonal chosen is up the main diagonal.

    • When offset is a negative integer, the diagonal chosen is down the main diagonal.

  • dim1 (int, optional) – Axis to be used as the first axis of the 2-D sub-arrays from which the diagonals should be taken. Default: 0 .

  • dim2 (int, optional) – Axis to be used as the second axis of the 2-D sub-arrays from which the diagonals should be taken. Default: 1 .

Returns

Tensor after embedding, has the same shape and dtype as input.

Raises

  • TypeError – If input or src is not a Tensor.

  • TypeError – If offset , dim1 or dim2 is not an integer.

Supported Platforms:

Ascend GPU CPU

Examples

>>> import mindspore as ms
>>> input = ms.ops.zeros((3,3))
>>> src = ms.ops.ones(2)
>>> out = ms.ops.diagonal_scatter(input, src, 1, dim1=1, dim2=0)
>>> print(out)
[[0. 0. 0.]
 [1. 0. 0.]
 [0. 1. 0.]]