ENH: add diag_indices, tril_indices, triu_indices (#692)

Co-authored-by: Quentin Barthélemy <q.barthelemy@gmail.com>
Co-authored-by: Lucas Colley <lucas.colley8@gmail.com>

Author: 3 people

docs/api-assorted.md

@@ -15,6 +15,7 @@
     cov
     create_diagonal
     default_dtype
+    diag_indices
     expand_dims
     isclose
     isin
@@ -27,6 +28,8 @@
     searchsorted
     setdiff1d
     sinc
+    tril_indices
+    triu_indices
     union1d
     unravel_index
 ```

src/array_api_extra/__init__.py

@@ -7,6 +7,7 @@
     broadcast_shapes,
     cov,
     create_diagonal,
+    diag_indices,
     expand_dims,
     isclose,
     isin,
@@ -18,6 +19,8 @@
     searchsorted,
     setdiff1d,
     sinc,
+    tril_indices,
+    triu_indices,
     union1d,
     unravel_index,
 )
@@ -44,6 +47,7 @@
     "cov",
     "create_diagonal",
     "default_dtype",
+    "diag_indices",
     "expand_dims",
     "isclose",
     "isin",
@@ -58,6 +62,8 @@
     "setdiff1d",
     "sinc",
     "testing",
+    "tril_indices",
+    "triu_indices",
     "union1d",
     "unravel_index",
 ]

src/array_api_extra/_delegation.py

@@ -21,13 +21,14 @@
     eager_shape,
     normalize_pad_width,
 )
-from ._lib._utils._typing import Array, DType
+from ._lib._utils._typing import Array, Device, DType
 
 __all__ = [
     "atleast_nd",
     "broadcast_shapes",
     "cov",
     "create_diagonal",
+    "diag_indices",
     "expand_dims",
     "isclose",
     "kron",
@@ -36,6 +37,8 @@
     "pad",
     "searchsorted",
     "sinc",
+    "tril_indices",
+    "triu_indices",
     "unravel_index",
 ]
 
@@ -310,6 +313,55 @@ def create_diagonal(
     return _funcs.create_diagonal(x, offset=offset, xp=xp)
 
 
+def diag_indices(
+    n: int, /, *, ndim: int = 2, device: Device | None = None, xp: ModuleType
+) -> tuple[Array, ...]:
+    """
+    Return the indices to access the main diagonal of an array.
+
+    Equivalent to :func:`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, optional
+        The number of dimensions. Default: ``2``.
+    device : Device, optional
+        The device on which to place the returned arrays. Default: current device.
+    xp : array_namespace
+        The standard-compatible namespace to create the indices in.
+
+    Returns
+    -------
+    tuple of array
+        1-D integer arrays of length ``n`` that together index
+        the main diagonal of an array of shape ``(n,) * ndim``.
+
+    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)
+    """
+    if n < 0:
+        msg = f"`n` must be non-negative, got {n}"
+        raise ValueError(msg)
+    if ndim < 1:
+        msg = f"`ndim` must be >= 1, got {ndim}"
+        raise ValueError(msg)
+    if device is None and (
+        is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp)
+    ):
+        return xp.diag_indices(n, ndim=ndim)
+    return _funcs.diag_indices(n, ndim=ndim, device=device, xp=xp)
+
+
 @deprecated(
     "`xpx.expand_dims` is deprecated and will be removed in v1.0.0. "
     "`xp.expand_dims` with support for a tuple of ints in `axis` "
@@ -802,11 +854,7 @@ def pad(
         # `torch/_numpy`'s implementation (avoids device transfers)
         pad_width_seq = normalize_pad_width(pad_width, x.ndim)
         # torch.nn.functional.pad counts dimensions from the last one
-        flat_pad_width = [
-            w
-            for pair in reversed(pad_width_seq)
-            for w in pair
-        ]
+        flat_pad_width = [w for pair in reversed(pad_width_seq) for w in pair]
         return xp.nn.functional.pad(x, tuple(flat_pad_width), value=constant_values)
 
     return _funcs.pad(x, pad_width, constant_values=constant_values, xp=xp)
@@ -1331,6 +1379,148 @@ def union1d(a: Array, b: Array, /, *, xp: ModuleType | None = None) -> Array:
     return _funcs.union1d(a, b, xp=xp)
 
 
+def tril_indices(
+    n: int,
+    /,
+    *,
+    offset: int = 0,
+    m: int | None = None,
+    device: Device | None = None,
+    xp: ModuleType,
+) -> tuple[Array, Array]:
+    """
+    Return the indices of the lower triangle of an ``(n, m)`` array.
+
+    Equivalent to :func:`numpy.tril_indices` with parameter ``k`` renamed to
+    ``offset`` to match :func:`array_api.linalg.diagonal`'s naming.
+
+    Parameters
+    ----------
+    n : int
+        The row dimension of the array.
+    offset : int, optional
+        Diagonal offset; ``0`` (default) is the main diagonal. Corresponds
+        to ``k`` in :func:`numpy.tril_indices`.
+    m : int, optional
+        The column dimension. If ``None`` (default), assumed equal to `n`.
+    device : Device, optional
+        The device on which to place the returned arrays. Default: current device.
+    xp : array_namespace
+        The standard-compatible namespace to create the indices in.
+
+    Returns
+    -------
+    tuple of array
+        Row and column indices ``(rows, cols)`` of the lower triangle of
+        the ``(n, m)`` matrix, shifted by `offset`.
+
+    Notes
+    -----
+    The generic fallback uses :func:`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.tril_indices(3, xp=xp)
+    >>> rows
+    Array([0, 1, 1, 2, 2, 2], dtype=array_api_strict.int64)
+    >>> cols
+    Array([0, 0, 1, 0, 1, 2], dtype=array_api_strict.int64)
+    """
+    if n < 0:
+        msg = f"`n` must be non-negative, got {n}"
+        raise ValueError(msg)
+    if m is not None and m < 0:
+        msg = f"`m` must be non-negative, got {m}"
+        raise ValueError(msg)
+    if device is None and (
+        is_numpy_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_jax_namespace(xp)
+        or is_dask_namespace(xp)
+    ):
+        return xp.tril_indices(n, k=offset, m=m)
+    if is_torch_namespace(xp):
+        # `torch.tril_indices` returns a 2xN tensor, not a tuple, and
+        # takes (row, col) rather than (n, *, m=None).
+        cols = n if m is None else m
+        idx = xp.tril_indices(n, cols, offset=offset, device=device)
+        return (idx[0], idx[1])
+    return _funcs.tril_indices(n, offset=offset, m=m, device=device, xp=xp)
+
+
+def triu_indices(
+    n: int,
+    /,
+    *,
+    offset: int = 0,
+    m: int | None = None,
+    device: Device | None = None,
+    xp: ModuleType,
+) -> tuple[Array, Array]:
+    """
+    Return the indices of the upper triangle of an ``(n, m)`` array.
+
+    Equivalent to :func:`numpy.triu_indices` with parameter ``k`` renamed to
+    ``offset`` to match :func:`array_api.linalg.diagonal`'s naming.
+
+    Parameters
+    ----------
+    n : int
+        The row dimension of the array.
+    offset : int, optional
+        Diagonal offset; ``0`` (default) is the main diagonal. Corresponds
+        to ``k`` in :func:`numpy.triu_indices`.
+    m : int, optional
+        The column dimension. If ``None`` (default), assumed equal to `n`.
+    device : Device, optional
+        The device on which to place the returned arrays. Default: current device.
+    xp : array_namespace
+        The standard-compatible namespace to create the indices in.
+
+    Returns
+    -------
+    tuple of array
+        Row and column indices ``(rows, cols)`` of the upper triangle of
+        the ``(n, m)`` matrix, shifted by `offset`.
+
+    Notes
+    -----
+    The generic fallback uses :func:`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)
+    """
+    if n < 0:
+        msg = f"`n` must be non-negative, got {n}"
+        raise ValueError(msg)
+    if m is not None and m < 0:
+        msg = f"`m` must be non-negative, got {m}"
+        raise ValueError(msg)
+    if device is None and (
+        is_numpy_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_jax_namespace(xp)
+        or is_dask_namespace(xp)
+    ):
+        return xp.triu_indices(n, k=offset, m=m)
+    if is_torch_namespace(xp):
+        cols = n if m is None else m
+        idx = xp.triu_indices(n, cols, offset=offset, device=device)
+        return (idx[0], idx[1])
+    return _funcs.triu_indices(n, offset=offset, m=m, device=device, xp=xp)
+
+
 def unravel_index(
     indices: Array,
     shape: tuple[int, ...],

src/array_api_extra/_lib/_funcs.py

@@ -30,13 +30,16 @@
     "broadcast_shapes",
     "cov",
     "create_diagonal",
+    "diag_indices",
     "expand_dims",
     "kron",
     "nunique",
     "pad",
     "searchsorted",
     "setdiff1d",
     "sinc",
+    "tril_indices",
+    "triu_indices",
 ]
 
 
@@ -312,6 +315,59 @@ def create_diagonal(
     return xp.reshape(diag, (*batch_dims, n, n))
 
 
+def diag_indices(
+    n: int, /, *, ndim: int, device: Device | None, xp: ModuleType
+) -> tuple[Array, ...]:  # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
+    idx = xp.arange(n, device=device)
+    return (idx,) * ndim
+
+
+def _tri_indices(
+    n: int,
+    *,
+    offset: int,
+    m: int | None,
+    upper: bool,
+    device: Device | None,
+    xp: ModuleType,
+) -> tuple[Array, Array]:  # numpydoc ignore=PR01,RT01
+    """Shared implementation for `tril_indices` and `triu_indices`."""
+    cols = n if m is None else m
+    rows = xp.arange(n, device=device)[:, xp.newaxis]
+    cols_a = xp.arange(cols, device=device)[xp.newaxis, :]
+    delta = cols_a - rows
+    mask = delta >= offset if upper else delta <= offset
+    r, c = xp.nonzero(mask)
+    return (r, c)
+
+
+def tril_indices(
+    n: int,
+    /,
+    *,
+    offset: int,
+    m: int | None,
+    device: Device | None,
+    xp: ModuleType,
+) -> tuple[Array, Array]:  # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
+    return _tri_indices(n, offset=offset, m=m, upper=False, device=device, xp=xp)
+
+
+def triu_indices(
+    n: int,
+    /,
+    *,
+    offset: int,
+    m: int | None,
+    device: Device | None,
+    xp: ModuleType,
+) -> tuple[Array, Array]:  # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
+    return _tri_indices(n, offset=offset, m=m, upper=True, device=device, xp=xp)
+
+
 def default_dtype(
     xp: ModuleType,
     kind: Literal[