add delegate func create_diagonal

Author: Enderdead

src/array_api_extra/__init__.py

@@ -4,6 +4,7 @@
     argpartition,
     atleast_nd,
     cov,
+    create_diagonal,
     expand_dims,
     isclose,
     isin,
@@ -17,7 +18,6 @@
 from ._lib._funcs import (
     apply_where,
     broadcast_shapes,
-    create_diagonal,
     default_dtype,
     kron,
     nunique,

src/array_api_extra/_delegation.py

@@ -19,7 +19,9 @@
 from ._lib._utils._typing import Array, DType
 
 __all__ = [
+    "atleast_nd",
     "cov",
+    "create_diagonal",
     "expand_dims",
     "isclose",
     "nan_to_num",
@@ -29,6 +31,55 @@
 ]
 
 
+def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array:
+    """
+    Recursively expand the dimension of an array to at least `ndim`.
+
+    Parameters
+    ----------
+    x : array
+        Input array.
+    ndim : int
+        The minimum number of dimensions for the result.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array
+        An array with ``res.ndim`` >= `ndim`.
+        If ``x.ndim`` >= `ndim`, `x` is returned.
+        If ``x.ndim`` < `ndim`, `x` is expanded by prepending new axes
+        until ``res.ndim`` equals `ndim`.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> x = xp.asarray([1])
+    >>> xpx.atleast_nd(x, ndim=3, xp=xp)
+    Array([[[1]]], dtype=array_api_strict.int64)
+
+    >>> x = xp.asarray([[[1, 2],
+    ...                  [3, 4]]])
+    >>> xpx.atleast_nd(x, ndim=1, xp=xp) is x
+    True
+    """
+    if xp is None:
+        xp = array_namespace(x)
+
+    if 1 <= ndim <= 3 and (
+        is_numpy_namespace(xp)
+        or is_jax_namespace(xp)
+        or is_dask_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_torch_namespace(xp)
+    ):
+        return getattr(xp, f"atleast_{ndim}d")(x)
+
+    return _funcs.atleast_nd(x, ndim=ndim, xp=xp)
+
+
 def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
     """
     Estimate a covariance matrix.
@@ -109,6 +160,69 @@ def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
     return _funcs.cov(m, xp=xp)
 
 
+def create_diagonal(
+    x: Array, /, *, offset: int = 0, xp: ModuleType | None = None
+) -> Array:
+    """
+    Construct a diagonal array.
+
+    Parameters
+    ----------
+    x : array
+        An array having shape ``(*batch_dims, k)``.
+    offset : int, optional
+        Offset from the leading diagonal (default is ``0``).
+        Use positive ints for diagonals above the leading diagonal,
+        and negative ints for diagonals below the leading diagonal.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array
+        An array having shape ``(*batch_dims, k+abs(offset), k+abs(offset))`` with `x`
+        on the diagonal (offset by `offset`).
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> x = xp.asarray([2, 4, 8])
+
+    >>> xpx.create_diagonal(x, xp=xp)
+    Array([[2, 0, 0],
+           [0, 4, 0],
+           [0, 0, 8]], dtype=array_api_strict.int64)
+
+    >>> xpx.create_diagonal(x, offset=-2, xp=xp)
+    Array([[0, 0, 0, 0, 0],
+           [0, 0, 0, 0, 0],
+           [2, 0, 0, 0, 0],
+           [0, 4, 0, 0, 0],
+           [0, 0, 8, 0, 0]], dtype=array_api_strict.int64)
+    """
+    if xp is None:
+        xp = array_namespace(x)
+
+    if x.ndim == 0:
+        err_msg = "`x` must be at least 1-dimensional."
+        raise ValueError(err_msg)
+
+    if is_torch_namespace(xp):
+        return xp.diag_embed(
+            atleast_nd(x, ndim=1, xp=xp), offset=offset, dim1=-2, dim2=-1
+        )
+
+    if (is_dask_namespace(xp) or is_cupy_namespace(xp)) and x.ndim < 2:
+        return xp.diag(x, k=offset)
+
+    if (is_jax_namespace(xp) or is_numpy_namespace(xp)) and x.ndim < 3:
+        batch_dim, n = eager_shape(x)[:-1], eager_shape(x, -1)[0] + abs(offset)
+        return xp.reshape(xp.diag(x, k=offset), (*batch_dim, n, n))
+
+    return _funcs.create_diagonal(x, offset=offset, xp=xp)
+
+
 def expand_dims(
     a: Array, /, *, axis: int | tuple[int, ...] = (0,), xp: ModuleType | None = None
 ) -> Array:
@@ -197,55 +311,6 @@ def expand_dims(
     return _funcs.expand_dims(a, axis=axis, xp=xp)
 
 
-def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array:
-    """
-    Recursively expand the dimension of an array to at least `ndim`.
-
-    Parameters
-    ----------
-    x : array
-        Input array.
-    ndim : int
-        The minimum number of dimensions for the result.
-    xp : array_namespace, optional
-        The standard-compatible namespace for `x`. Default: infer.
-
-    Returns
-    -------
-    array
-        An array with ``res.ndim`` >= `ndim`.
-        If ``x.ndim`` >= `ndim`, `x` is returned.
-        If ``x.ndim`` < `ndim`, `x` is expanded by prepending new axes
-        until ``res.ndim`` equals `ndim`.
-
-    Examples
-    --------
-    >>> import array_api_strict as xp
-    >>> import array_api_extra as xpx
-    >>> x = xp.asarray([1])
-    >>> xpx.atleast_nd(x, ndim=3, xp=xp)
-    Array([[[1]]], dtype=array_api_strict.int64)
-
-    >>> x = xp.asarray([[[1, 2],
-    ...                  [3, 4]]])
-    >>> xpx.atleast_nd(x, ndim=1, xp=xp) is x
-    True
-    """
-    if xp is None:
-        xp = array_namespace(x)
-
-    if 1 <= ndim <= 3 and (
-        is_numpy_namespace(xp)
-        or is_jax_namespace(xp)
-        or is_dask_namespace(xp)
-        or is_cupy_namespace(xp)
-        or is_torch_namespace(xp)
-    ):
-        return getattr(xp, f"atleast_{ndim}d")(x)
-
-    return _funcs.atleast_nd(x, ndim=ndim, xp=xp)
-
-
 def isclose(
     a: Array | complex,
     b: Array | complex,

src/array_api_extra/_lib/_funcs.py

@@ -295,53 +295,9 @@ def one_hot(
 
 
 def create_diagonal(
-    x: Array, /, *, offset: int = 0, xp: ModuleType | None = None
-) -> Array:
-    """
-    Construct a diagonal array.
-
-    Parameters
-    ----------
-    x : array
-        An array having shape ``(*batch_dims, k)``.
-    offset : int, optional
-        Offset from the leading diagonal (default is ``0``).
-        Use positive ints for diagonals above the leading diagonal,
-        and negative ints for diagonals below the leading diagonal.
-    xp : array_namespace, optional
-        The standard-compatible namespace for `x`. Default: infer.
-
-    Returns
-    -------
-    array
-        An array having shape ``(*batch_dims, k+abs(offset), k+abs(offset))`` with `x`
-        on the diagonal (offset by `offset`).
-
-    Examples
-    --------
-    >>> import array_api_strict as xp
-    >>> import array_api_extra as xpx
-    >>> x = xp.asarray([2, 4, 8])
-
-    >>> xpx.create_diagonal(x, xp=xp)
-    Array([[2, 0, 0],
-           [0, 4, 0],
-           [0, 0, 8]], dtype=array_api_strict.int64)
-
-    >>> xpx.create_diagonal(x, offset=-2, xp=xp)
-    Array([[0, 0, 0, 0, 0],
-           [0, 0, 0, 0, 0],
-           [2, 0, 0, 0, 0],
-           [0, 4, 0, 0, 0],
-           [0, 0, 8, 0, 0]], dtype=array_api_strict.int64)
-    """
-    if xp is None:
-        xp = array_namespace(x)
-
-    if x.ndim == 0:
-        err_msg = "`x` must be at least 1-dimensional."
-        raise ValueError(err_msg)
-
+    x: Array, /, *, offset: int = 0, xp: ModuleType
+) -> Array:  # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
     x_shape = eager_shape(x)
     batch_dims = x_shape[:-1]
     n = x_shape[-1] + abs(offset)