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 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)