jax.numpy.triu_indices_from

jax.numpy.triu_indices_from(arr, k=0)

Return the indices of upper triangle of a given array.

JAX implementation of numpy.triu_indices_from().

Parameters:

• arr (ArrayLike | SupportsShape) – input array. Must have arr.ndim == 2.

• k (int) – optional, int, default=0. Specifies the sub-diagonal on and above which the indices of upper triangle are returned. k=0 refers to main diagonal, k<0 refers to sub-diagonal below the main diagonal and k>0 refers to sub-diagonal above the main diagonal.

Returns:

A tuple of two arrays containing the indices of the upper triangle, one along each axis.

Return type:

tuple[Array, Array]

See also

• jax.numpy.tril_indices_from(): Returns the indices of lower triangle of a given array.

• jax.numpy.triu_indices(): Returns the indices of upper triangle of an array of size (n, m).

• jax.numpy.triu(): Return an upper triangle of an array.

Examples

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

Elements indexed by jnp.triu_indices_from correspond to those in the output of jnp.triu.

>>> ind = jnp.triu_indices_from(arr)
>>> arr[ind]
Array([1, 2, 3, 5, 6, 9], dtype=int32)
>>> jnp.triu(arr)
Array([[1, 2, 3],
       [0, 5, 6],
       [0, 0, 9]], dtype=int32)

When k > 0:

>>> jnp.triu_indices_from(arr, k=1)
(Array([0, 0, 1], dtype=int32), Array([1, 2, 2], dtype=int32))

When k < 0:

>>> jnp.triu_indices_from(arr, k=-1)
(Array([0, 0, 0, 1, 1, 1, 2, 2], dtype=int32), Array([0, 1, 2, 0, 1, 2, 1, 2], dtype=int32))