mindspore.numpy.diag_indices

mindspore.numpy.diag_indices(n, ndim=2)

Returns the indices to access the main diagonal of an array.

This returns a tuple of indices that can be used to access the main diagonal of an array a with a.ndim >= 2 dimensions and shape (n, n, …, n). For a.ndim = 2 this is the usual diagonal, for a.ndim > 2 this is the set of indices to access a[i, i, ..., i] for i = [0..n-1].

参数

• n (int) – The size, along each dimension, of the arrays for which the returned indices can be used.

• ndim (int, optional) – The number of dimensions.

返回

Tuple of Tensor.

异常

TypeError – If input are not integers.

Supported Platforms:

Ascend GPU CPU

样例

>>> import mindspore.numpy as np
>>> output = np.diag_indices(5, 3)
>>> print(output)
(Tensor(shape=[5], dtype=Int32, value= [0, 1, 2, 3, 4]),
Tensor(shape=[5], dtype=Int32, value= [0, 1, 2, 3, 4]),
Tensor(shape=[5], dtype=Int32, value= [0, 1, 2, 3, 4]))