triu_indices_from

class brainunit.math.triu_indices_from(arr, k=0)

Return the indices for the upper-triangle of an (n, m) array.

Parameters:

• arr (saiunit.Quantity | Array | ndarray | bool | number | bool | int | float | complex) – The array for which the returned indices will be valid.

• k (int | None) – Diagonal offset. k = 0 is the main diagonal, k < 0 is below, and k > 0 is above.

Returns:

out – Row and column indices for the upper triangle.

Return type:

Tuple[Array, Array]

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> row, col = u.math.triu_indices_from(jnp.ones((3, 3)))
>>> row
Array([0, 0, 0, 1, 1, 2], dtype=int32)