TST/BUG: run all tests on all backends; fix backend-specific bugs

Author: crusaderky

.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

@@ -15,3 +15,15 @@
     setdiff1d
     sinc
 ```
+
+## Test tools
+
+```{eval-rst}
+.. currentmodule:: array_api_extra.testing
+.. autosummary::
+    :nosignatures:
+    :toctree: generated
+
+    xp_assert_equal
+    xp_assert_close
+```

pixi.lock

@@ -180,7 +180,7 @@ xfail_strict = true
 filterwarnings = ["error"]
 log_cli_level = "INFO"
 testpaths = ["tests"]
-
+markers = ["skip_xp_backend(library, *, reason=None): Skip test for a specific backend"]
 
 # Coverage
 

pyproject.toml

@@ -214,8 +214,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))
 
 
@@ -407,9 +411,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(
@@ -632,8 +635,7 @@ def pad(
         dtype=x.dtype,
         device=_compat.device(x),
     )
-    padded[tuple(slices)] = x
-    return padded
+    return at(padded, tuple(slices)).set(x)
 
 
 class _AtOp(Enum):

src/array_api_extra/_funcs.py

@@ -6,20 +6,35 @@
     from ..._array_api_compat_vendor import (  # pyright: ignore[reportMissingImports]
         array_namespace,
         device,
+        is_cupy_namespace,
         is_jax_array,
+        is_jax_namespace,
+        is_pydata_sparse_namespace,
+        is_torch_namespace,
         is_writeable_array,
+        size,
     )
 except ImportError:
     from array_api_compat import (  # pyright: ignore[reportMissingTypeStubs]
         array_namespace,
         device,
+        is_cupy_namespace,
         is_jax_array,
+        is_jax_namespace,
+        is_pydata_sparse_namespace,
+        is_torch_namespace,
         is_writeable_array,
+        size,
     )
 
 __all__ = [
     "array_namespace",
     "device",
+    "is_cupy_namespace",
     "is_jax_array",
+    "is_jax_namespace",
+    "is_pydata_sparse_namespace",
+    "is_torch_namespace",
     "is_writeable_array",
+    "size",
 ]

src/array_api_extra/_lib/_compat.py

@@ -18,5 +18,10 @@ def array_namespace(
     use_compat: bool | None = None,
 ) -> ArrayModule: ...
 def device(x: Array, /) -> Device: ...
+def is_cupy_namespace(x: object, /) -> bool: ...
 def is_jax_array(x: object, /) -> bool: ...
+def is_jax_namespace(x: object, /) -> bool: ...
+def is_pydata_sparse_namespace(x: object, /) -> bool: ...
+def is_torch_namespace(x: object, /) -> bool: ...
 def is_writeable_array(x: object, /) -> bool: ...
+def size(x: Array, /) -> int | None: ...

src/array_api_extra/_lib/_compat.pyi

@@ -54,7 +54,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])

src/array_api_extra/_lib/_utils.py

@@ -0,0 +1,130 @@
+"""Testing utilities."""
+
+from ._lib._compat import (
+    array_namespace,
+    is_cupy_namespace,
+    is_pydata_sparse_namespace,
+    is_torch_namespace,
+)
+from ._lib._typing import Array
+
+__all__ = ["xp_assert_close", "xp_assert_equal"]
+
+
+def _check_shape_dtype(actual: Array, desired: Array) -> None:
+    """
+    Assert that 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).
+    """
+    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}".format(
+        actual.dtype, desired.dtype
+    )
+    assert actual.dtype == desired.dtype, msg
+
+
+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 = array_namespace(actual, desired)
+    _check_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 = array_namespace(actual, desired)
+    _check_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/testing.py

@@ -0,0 +1 @@
+"""Needed to import .conftest from the test modules."""

tests/__init__.py

@@ -0,0 +1,86 @@
+"""Pytest fixtures."""
+
+from enum import Enum
+from typing import cast
+
+import pytest
+
+from array_api_extra._lib._compat import array_namespace
+from array_api_extra._lib._compat import device as get_device
+from array_api_extra._lib._typing import Device, ModuleType
+
+
+class Library(Enum):
+    """All array libraries explicitly tested by array-api-extra."""
+
+    ARRAY_API_STRICT = "array_api_strict"
+    NUMPY = "numpy"
+    NUMPY_READONLY = "numpy_readonly"
+    CUPY = "cupy"
+    TORCH = "torch"
+    DASK_ARRAY = "dask.array"
+    SPARSE = "sparse"
+    JAX_NUMPY = "jax.numpy"
+
+    def __str__(self) -> str:  # type: ignore[explicit-override]  # pyright: ignore[reportImplicitOverride]  # numpydoc ignore=RT01
+        """Pretty-print parameterized test names."""
+        return self.value
+
+
+@pytest.fixture(params=tuple(Library))
+def library(request: pytest.FixtureRequest) -> Library:  # numpydoc ignore=PR01,RT03
+    """
+    Parameterized fixture that iterates on all libraries.
+
+    Returns
+    -------
+    The current Library enum.
+    """
+    elem = cast(Library, request.param)
+
+    for marker in request.node.iter_markers("skip_xp_backend"):
+        skip_library = marker.kwargs.get("library") or marker.args[0]  # type: ignore[no-untyped-usage]
+        if not isinstance(skip_library, Library):
+            msg = "argument of skip_xp_backend must be a Library enum"
+            raise TypeError(msg)
+        if skip_library == elem:
+            reason = cast(str, marker.kwargs.get("reason", "skip_xp_backend"))
+            pytest.skip(reason=reason)
+
+    return elem
+
+
+@pytest.fixture
+def xp(library: Library) -> ModuleType:  # numpydoc ignore=PR01,RT03
+    """
+    Parameterized fixture that iterates on all libraries.
+
+    Returns
+    -------
+    The current array namespace.
+    """
+    name = "numpy" if library == Library.NUMPY_READONLY else library.value
+    xp = pytest.importorskip(name)
+    if library == Library.JAX_NUMPY:
+        import jax  # type: ignore[import-not-found]  # pyright: ignore[reportMissingImports]
+
+        jax.config.update("jax_enable_x64", True)
+
+    # Possibly wrap module with array_api_compat
+    return array_namespace(xp.empty(0))
+
+
+@pytest.fixture
+def device(
+    library: Library, xp: ModuleType
+) -> Device:  # numpydoc ignore=PR01,RT01,RT03
+    """
+    Return a valid device for the backend.
+
+    Where possible, return a device that is not the default one.
+    """
+    if library == Library.ARRAY_API_STRICT:
+        d = xp.Device("device1")
+        assert get_device(xp.empty(0)) != d
+        return d
+    return get_device(xp.empty(0))