jax.numpy.diag_indices_from

jax.numpy.diag_indices_from(arr)

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

JAX implementation of numpy.diag_indices_from().

Parameters:

arr (ArrayLike) – Input array. Must be at least 2-dimensional and have equal length along all dimensions.

Returns:

A tuple of arrays containing the indices to access the main diagonal of the input array.

Return type:

tuple[Array, …]

See also

• jax.numpy.diag_indices()

• jax.numpy.diagonal()

Examples

>>> arr = jnp.array([[1, 2, 3],
...                  [4, 5, 6],
...                  [7, 8, 9]])
>>> jnp.diag_indices_from(arr)
(Array([0, 1, 2], dtype=int32), Array([0, 1, 2], dtype=int32))
>>> arr = jnp.array([[[1, 2], [3, 4]],
...                  [[5, 6], [7, 8]]])
>>> jnp.diag_indices_from(arr)
(Array([0, 1], dtype=int32),
Array([0, 1], dtype=int32),
Array([0, 1], dtype=int32))