triu_indices

array_api_extra.triu_indices(n, /, *, offset=0, m=None, device=None, xp)

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

Equivalent to numpy.triu_indices() with parameter k renamed to offset to match array_api.linalg.diagonal()’s naming.

Parameters:

  • n (int) – The row dimension of the array.

  • offset (int) – Diagonal offset; 0 (default) is the main diagonal. Corresponds to k in numpy.triu_indices().

  • m (int | None) – The column dimension. If None (default), assumed equal to n.

  • device (object | None) – The device on which to place the returned arrays. Default: current device.

  • xp (ModuleType) – The standard-compatible namespace to create the indices in.

Returns:

Row and column indices (rows, cols) of the upper triangle of the (n, m) matrix, shifted by offset.

Return type:

tuple[object, object]

Notes

The generic fallback uses array_api.nonzero(), so namespaces without nonzero are not supported on that path.

Examples

>>> import array_api_strict as xp
>>> import array_api_extra as xpx
>>> rows, cols = xpx.triu_indices(3, xp=xp)
>>> rows
Array([0, 0, 0, 1, 1, 2], dtype=array_api_strict.int64)
>>> cols
Array([0, 1, 2, 1, 2, 2], dtype=array_api_strict.int64)