mindspore.numpy.trace

mindspore.numpy.trace(a, offset=0, axis1=0, axis2=1, dtype=None)[source]

Returns the sum along diagonals of the array.

If a is 2-D, the sum along its diagonal with the given offset is returned, i.e., the sum of elements a[i,i+offset] for all i. If a has more than two dimensions, then the axes specified by axis1 and axis2 are used to determine the 2-D sub-arrays whose traces are returned. The shape of the resulting array is the same as that of a with axis1 and axis2 removed.

Note

On GPU, the supported dtypes are np.float16, and np.float32. On CPU, the supported dtypes are np.float16, np.float32, and np.float64.

Parameters

  • a (Tensor) – Array from which the diagonals are taken.

  • offset (int, optional) – Offset of the diagonal from the main diagonal. Can be positive or negative. Defaults to main diagonal.

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

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

  • dtype (mindspore.dtype, optional) – defaults to None. Overrides the dtype of the output Tensor.

Returns

Tensor, sum_along_diagonals. If a is 2-D, the sum along the diagonal is returned. If a has larger dimensions, then an array of sums along diagonals is returned.

Raises

ValueError – if the input tensor has less than two dimensions.

Supported Platforms:

Ascend GPU CPU

Examples

>>> import mindspore.numpy as np
>>> output = np.trace(np.eye(3))
>>> print(output)
3.0
>>> a = np.arange(8).reshape((2,2,2))
>>> output = np.trace(a)
>>> print(output)
[6 8]
>>> a = np.arange(24).reshape((2,2,2,3))
>>> output = np.trace(a).shape
>>> print(output)
(2, 3)