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