diag_indices

array_api_extra.diag_indices(n, /, *, ndim=2, device=None, xp)

Return the indices to access the main diagonal of an array.

Equivalent to numpy.diag_indices().

Parameters:

  • n (int) – The size of each dimension of the (hyper-)cube (n, n, ..., n) that the returned indices index into.

  • ndim (int) – The number of dimensions. Default: 2.

  • 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:

1-D integer arrays of length n that together index the main diagonal of an array of shape (n,) * ndim.

Return type:

tuple[object, ...]

Examples

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