jax.numpy.triu_indices_from#

jax.numpy.triu_indices_from(arr, k=0)[source]#

Return the indices for the upper-triangle of arr.

LAX-backend implementation of numpy.triu_indices_from().

Original docstring below.

See triu_indices for full details.

Parameters:

  • arr (ndarray, shape(N, N)) – The indices will be valid for square arrays.

  • k (int, optional) – Diagonal offset (see triu for details).

Returns:

triu_indices_from – Indices for the upper-triangle of arr.

Return type:

tuple, shape(2) of ndarray, shape(N)