jax.numpy.diag_indices#

jax.numpy.diag_indices(n,ndim=2)[source]#

Return indices for accessing the main diagonal of a multidimensional array.

JAX implementation ofnumpy.diag_indices().

Parameters:

• n (int) – int. The size of each dimension of the square array.

• ndim (int) – optional, int, default=2. The number of dimensions of the array.

Returns:

A tuple of arrays, each of lengthn, containing the indices to accessthe main diagonal.

Return type:

tuple[Array, …]

See also

• jax.numpy.diag_indices_from()

• jax.numpy.diagonal()

Examples

>>>jnp.diag_indices(3)(Array([0, 1, 2], dtype=int32), Array([0, 1, 2], dtype=int32))>>>jnp.diag_indices(4,ndim=3)(Array([0, 1, 2, 3], dtype=int32),Array([0, 1, 2, 3], dtype=int32),Array([0, 1, 2, 3], dtype=int32))