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)