Merge branch 'main' into pad-delegate

Author: lucascolley

.pre-commit-config.yaml

@@ -44,9 +44,9 @@ repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: "v0.8.2"
     hooks:
+      - id: ruff-format
       - id: ruff
         args: ["--fix", "--show-fixes"]
-      - id: ruff-format
 
   - repo: https://github.com/codespell-project/codespell
     rev: "v2.3.0"

docs/api-reference.md

@@ -12,6 +12,7 @@
     create_diagonal
     expand_dims
     kron
+    nunique
     pad
     setdiff1d
     sinc

pixi.lock

@@ -184,7 +184,7 @@ filterwarnings = [
 ]
 log_cli_level = "INFO"
 testpaths = ["tests"]
-
+markers = ["skip_xp_backend(library, *, reason=None): Skip test for a specific backend"]
 
 # Coverage
 
@@ -303,6 +303,7 @@ messages_control.disable = [
     "line-too-long",
     "missing-module-docstring",
     "missing-function-docstring",
+    "too-many-lines",
     "wrong-import-position",
 ]
 
@@ -319,6 +320,7 @@ checks = [
 exclude = [ # don't report on objects that match any of these regex
     '.*test_at.*',
     '.*test_funcs.*',
+    '.*test_testing.*',
     '.*test_utils.*',
     '.*test_version.*',
     '.*test_vendor.*',

pyproject.toml

@@ -8,6 +8,7 @@
     create_diagonal,
     expand_dims,
     kron,
+    nunique,
     setdiff1d,
     sinc,
 )
@@ -23,6 +24,7 @@
     "create_diagonal",
     "expand_dims",
     "kron",
+    "nunique",
     "pad",
     "setdiff1d",
     "sinc",

src/array_api_extra/__init__.py

@@ -3,6 +3,7 @@
 # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
 from __future__ import annotations
 
+import math
 import operator
 import warnings
 from collections.abc import Callable
@@ -21,6 +22,7 @@
     "create_diagonal",
     "expand_dims",
     "kron",
+    "nunique",
     "pad",
     "setdiff1d",
     "sinc",
@@ -210,8 +212,12 @@ def create_diagonal(
         raise ValueError(err_msg)
     n = x.shape[0] + abs(offset)
     diag = xp.zeros(n**2, dtype=x.dtype, device=_compat.device(x))
-    i = offset if offset >= 0 else abs(offset) * n
-    diag[i : min(n * (n - offset), diag.shape[0]) : n + 1] = x
+
+    start = offset if offset >= 0 else abs(offset) * n
+    stop = min(n * (n - offset), diag.shape[0])
+    step = n + 1
+    diag = at(diag)[start:stop:step].set(x)
+
     return xp.reshape(diag, (n, n))
 
 
@@ -403,9 +409,8 @@ def kron(a: Array, b: Array, /, *, xp: ModuleType | None = None) -> Array:
     result = xp.multiply(a_arr, b_arr)
 
     # Reshape back and return
-    a_shape = xp.asarray(a_shape)
-    b_shape = xp.asarray(b_shape)
-    return xp.reshape(result, tuple(xp.multiply(a_shape, b_shape)))
+    res_shape = tuple(a_s * b_s for a_s, b_s in zip(a_shape, b_shape, strict=True))
+    return xp.reshape(result, res_shape)
 
 
 def setdiff1d(
@@ -565,7 +570,7 @@ def pad(
     if isinstance(pad_width, tuple):
         pad_width = [pad_width] * x_ndim
 
-    # https://github.com/data-apis/array-api-extra/pull/82#discussion_r1905688819
+    # https://github.com/python/typeshed/issues/13376
     slices: list[slice] = []  # type: ignore[no-any-explicit]
     newshape: list[int] = []
     for ax, w_tpl in enumerate(pad_width):
@@ -592,8 +597,43 @@ def pad(
         dtype=x.dtype,
         device=_compat.device(x),
     )
-    padded[tuple(slices)] = x
-    return padded
+    return at(padded, tuple(slices)).set(x)
+
+
+def nunique(x: Array, /, *, xp: ModuleType | None = None) -> Array:
+    """
+    Count the number of unique elements in an array.
+
+    Compatible with JAX and Dask, whose laziness would be otherwise
+    problematic.
+
+    Parameters
+    ----------
+    x : Array
+        Input array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array: 0-dimensional integer array
+        The number of unique elements in `x`. It can be lazy.
+    """
+    if xp is None:
+        xp = array_namespace(x)
+
+    if is_jax_array(x):
+        # size= is JAX-specific
+        # https://github.com/data-apis/array-api/issues/883
+        _, counts = xp.unique_counts(x, size=_compat.size(x))
+        return xp.astype(counts, xp.bool).sum()
+
+    _, counts = xp.unique_counts(x)
+    n = _compat.size(counts)
+    # FIXME https://github.com/data-apis/array-api-compat/pull/231
+    if n is None or math.isnan(n):  # e.g. Dask, ndonnx
+        return xp.astype(counts, xp.bool).sum()
+    return xp.asarray(n, device=_compat.device(x))
 
 
 class _AtOp(Enum):

src/array_api_extra/_lib/_funcs.py

@@ -0,0 +1,146 @@
+"""
+Testing utilities.
+
+Note that this is private API; don't expect it to be stable.
+"""
+
+from types import ModuleType
+
+from ._utils._compat import (
+    array_namespace,
+    is_cupy_namespace,
+    is_pydata_sparse_namespace,
+    is_torch_namespace,
+)
+from ._utils._typing import Array
+
+__all__ = ["xp_assert_close", "xp_assert_equal"]
+
+
+def _check_ns_shape_dtype(
+    actual: Array, desired: Array
+) -> ModuleType:  # numpydoc ignore=RT03
+    """
+    Assert that namespace, shape and dtype of the two arrays match.
+
+    Parameters
+    ----------
+    actual : Array
+        The array produced by the tested function.
+    desired : Array
+        The expected array (typically hardcoded).
+
+    Returns
+    -------
+    Arrays namespace.
+    """
+    actual_xp = array_namespace(actual)  # Raises on scalars and lists
+    desired_xp = array_namespace(desired)
+
+    msg = f"namespaces do not match: {actual_xp} != f{desired_xp}"
+    assert actual_xp == desired_xp, msg
+
+    msg = f"shapes do not match: {actual.shape} != f{desired.shape}"
+    assert actual.shape == desired.shape, msg
+
+    msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}"
+    assert actual.dtype == desired.dtype, msg
+
+    return desired_xp
+
+
+def xp_assert_equal(actual: Array, desired: Array, err_msg: str = "") -> None:
+    """
+    Array-API compatible version of `np.testing.assert_array_equal`.
+
+    Parameters
+    ----------
+    actual : Array
+        The array produced by the tested function.
+    desired : Array
+        The expected array (typically hardcoded).
+    err_msg : str, optional
+        Error message to display on failure.
+    """
+    xp = _check_ns_shape_dtype(actual, desired)
+
+    if is_cupy_namespace(xp):
+        xp.testing.assert_array_equal(actual, desired, err_msg=err_msg)
+    elif is_torch_namespace(xp):
+        # PyTorch recommends using `rtol=0, atol=0` like this
+        # to test for exact equality
+        xp.testing.assert_close(
+            actual,
+            desired,
+            rtol=0,
+            atol=0,
+            equal_nan=True,
+            check_dtype=False,
+            msg=err_msg or None,
+        )
+    else:
+        import numpy as np  # pylint: disable=import-outside-toplevel
+
+        if is_pydata_sparse_namespace(xp):
+            actual = actual.todense()
+            desired = desired.todense()
+
+        # JAX uses `np.testing`
+        np.testing.assert_array_equal(actual, desired, err_msg=err_msg)
+
+
+def xp_assert_close(
+    actual: Array,
+    desired: Array,
+    *,
+    rtol: float | None = None,
+    atol: float = 0,
+    err_msg: str = "",
+) -> None:
+    """
+    Array-API compatible version of `np.testing.assert_allclose`.
+
+    Parameters
+    ----------
+    actual : Array
+        The array produced by the tested function.
+    desired : Array
+        The expected array (typically hardcoded).
+    rtol : float, optional
+        Relative tolerance. Default: dtype-dependent.
+    atol : float, optional
+        Absolute tolerance. Default: 0.
+    err_msg : str, optional
+        Error message to display on failure.
+    """
+    xp = _check_ns_shape_dtype(actual, desired)
+
+    floating = xp.isdtype(actual.dtype, ("real floating", "complex floating"))
+    if rtol is None and floating:
+        # multiplier of 4 is used as for `np.float64` this puts the default `rtol`
+        # roughly half way between sqrt(eps) and the default for
+        # `numpy.testing.assert_allclose`, 1e-7
+        rtol = xp.finfo(actual.dtype).eps ** 0.5 * 4
+    elif rtol is None:
+        rtol = 1e-7
+
+    if is_cupy_namespace(xp):
+        xp.testing.assert_allclose(
+            actual, desired, rtol=rtol, atol=atol, err_msg=err_msg
+        )
+    elif is_torch_namespace(xp):
+        xp.testing.assert_close(
+            actual, desired, rtol=rtol, atol=atol, equal_nan=True, msg=err_msg or None
+        )
+    else:
+        import numpy as np  # pylint: disable=import-outside-toplevel
+
+        if is_pydata_sparse_namespace(xp):
+            actual = actual.to_dense()
+            desired = desired.to_dense()
+
+        # JAX uses `np.testing`
+        assert isinstance(rtol, float)
+        np.testing.assert_allclose(
+            actual, desired, rtol=rtol, atol=atol, err_msg=err_msg
+        )

src/array_api_extra/_lib/_testing.py

@@ -10,8 +10,10 @@
         is_jax_array,
         is_jax_namespace,
         is_numpy_namespace,
+        is_pydata_sparse_namespace,
         is_torch_namespace,
         is_writeable_array,
+        size,
     )
 except ImportError:
     from array_api_compat import (  # pyright: ignore[reportMissingTypeStubs]
@@ -21,8 +23,10 @@
         is_jax_array,
         is_jax_namespace,
         is_numpy_namespace,
+        is_pydata_sparse_namespace,
         is_torch_namespace,
         is_writeable_array,
+        size,
     )
 
 __all__ = [
@@ -32,6 +36,8 @@
     "is_jax_array",
     "is_jax_namespace",
     "is_numpy_namespace",
+    "is_pydata_sparse_namespace",
     "is_torch_namespace",
     "is_writeable_array",
+    "size",
 ]

src/array_api_extra/_lib/_utils/_compat.py

@@ -23,4 +23,6 @@ def is_jax_namespace(xp: ModuleType, /) -> bool: ...
 def is_numpy_namespace(xp: ModuleType, /) -> bool: ...
 def is_torch_namespace(xp: ModuleType, /) -> bool: ...
 def is_jax_array(x: object, /) -> bool: ...
+def is_pydata_sparse_namespace(xp: ModuleType, /) -> bool: ...
 def is_writeable_array(x: object, /) -> bool: ...
+def size(x: Array, /) -> int | None: ...

src/array_api_extra/_lib/_utils/_compat.pyi

@@ -56,7 +56,9 @@ def in1d(
     order = xp.argsort(ar, stable=True)
     reverse_order = xp.argsort(order, stable=True)
     sar = xp.take(ar, order, axis=0)
-    if sar.size >= 1:
+    ar_size = _compat.size(sar)
+    assert ar_size is not None, "xp.unique*() on lazy backends raises"
+    if ar_size >= 1:
         bool_ar = sar[1:] != sar[:-1] if invert else sar[1:] == sar[:-1]
     else:
         bool_ar = xp.asarray([False]) if invert else xp.asarray([True])