mindspore.mint.trace

View Source On AtomGit
mindspore.mint.trace(input)[source]

Return the sum of the elements along the diagonal of the input tensor.

Parameters:

input (Tensor) – 2-D input tensor.

Returns:

Tensor, when the data type of input is integer or bool, its data type is mindspore.int64, otherwise it is the same as input, and size equals to 1.

Raises:

ValueError – If the dimension of input is not equal to 2.

Supported Platforms:

Ascend

Examples

>>> import mindspore
>>> input = mindspore.tensor(mindspore.mint.arange(1, 13).reshape(3, 4), mindspore.float32)
>>> print(input)
[[ 1.  2.  3.  4.]
 [ 5.  6.  7.  8.]
 [ 9. 10. 11. 12.]]
>>> output = mindspore.mint.trace(input)
>>> print(output)
18.0