mindspore.nn.MatrixSetDiag

class mindspore.nn.MatrixSetDiag

Modifies the batched diagonal part of a batched tensor.

Assume x has \(k+1\) dimensions \([I, J, K, ..., M, N]\) and diagonal has \(k\) dimensions \([I, J, K, ..., min(M, N)]\). Then the output is a tensor of rank \(k+1\) with dimensions \([I, J, K, ..., M, N]\) where:

\[output[i, j, k, ..., m, n] = diagnoal[i, j, k, ..., n]\ for\ m == n\]

\[output[i, j, k, ..., m, n] = x[i, j, k, ..., m, n]\ for\ m != n\]

Inputs:

• x (Tensor) - The batched tensor. Rank k+1, where k >= 1. It can be one of the following data types: float32, float16, int32, int8, and uint8.

• diagonal (Tensor) - The diagonal values. Must have the same type as input x. Rank k, where k >= 1.

Outputs:

Tensor, has the same type and shape as input x.

Raises

• TypeError – If dtype of x or diagonal is not one of float32, float16, int32, int8 or uint8.

• ValueError – If length of shape of x is less than 2.

• ValueError – If x_shape[-2] < x_shape[-1] and x_shape[:-1] != diagonal_shape.

• ValueError – If x_shape[-2] >= x_shape[-1] and x_shape[:-2] + x_shape[-1:] != diagonal_shape.

Supported Platforms:

Ascend

Examples

>>> x = Tensor([[[-1, 0], [0, 1]], [[-1, 0], [0, 1]], [[-1, 0], [0, 1]]], mindspore.float32)
>>> diagonal = Tensor([[-1., 2.], [-1., 1.], [-1., 1.]], mindspore.float32)
>>> matrix_set_diag = nn.MatrixSetDiag()
>>> output = matrix_set_diag(x, diagonal)
>>> print(output)
[[[-1.  0.]
  [ 0.  2.]]
 [[-1.  0.]
  [ 0.  1.]]
 [[-1.  0.]
  [ 0.  1.]]]