jax.numpy.triu_indices_from#
- jax.numpy.triu_indices_from(arr, k=0)[source]#
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 andk>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:
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 ofjnp.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))