Merge branch 'main' into apply

Author: crusaderky

pixi.lock

@@ -125,22 +125,21 @@ python = "~=3.10.0"
 python = "~=3.13.0"
 
 # Backends that can run on CPU-only hosts
-[tool.pixi.feature.backends.target.linux-64.dependencies]
+[tool.pixi.feature.backends.dependencies]
 pytorch = "*"
 dask = "*"
-sparse = ">=0.15"
+numba = "*"  # sparse dependency
+
+[tool.pixi.feature.backends.pypi-dependencies]
+sparse = { version = ">= 0.16.0b3" }
+
+[tool.pixi.feature.backends.target.linux-64.dependencies]
 jax = "*"
 
 [tool.pixi.feature.backends.target.osx-arm64.dependencies]
-pytorch = "*"
-dask = "*"
-sparse = ">=0.15"
 jax = "*"
 
 [tool.pixi.feature.backends.target.win-64.dependencies]
-pytorch = "*"
-dask = "*"
-sparse = ">=0.15"
 # jax = "*"  # unavailable
 
 # Backends that require a GPU host and a CUDA driver

pyproject.toml

@@ -125,7 +125,7 @@ def pad(
     pad_width: int | tuple[int, int] | Sequence[tuple[int, int]],
     mode: Literal["constant"] = "constant",
     *,
-    constant_values: bool | int | float | complex = 0,
+    constant_values: complex = 0,
     xp: ModuleType | None = None,
 ) -> Array:
     """
@@ -168,7 +168,7 @@ def pad(
         pad_width = xp.flip(pad_width, axis=(0,)).flatten()
         return xp.nn.functional.pad(x, tuple(pad_width), value=constant_values)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 
-    if _delegate(xp, Backend.NUMPY, Backend.JAX, Backend.CUPY):
+    if _delegate(xp, Backend.NUMPY, Backend.JAX, Backend.CUPY, Backend.SPARSE):
         return xp.pad(x, pad_width, mode, constant_values=constant_values)
 
     return _funcs.pad(x, pad_width, constant_values=constant_values, xp=xp)

src/array_api_extra/_delegation.py

@@ -575,7 +575,7 @@ def pad(
     x: Array,
     pad_width: int | tuple[int, int] | Sequence[tuple[int, int]],
     *,
-    constant_values: bool | int | float | complex = 0,
+    constant_values: complex = 0,
     xp: ModuleType,
 ) -> Array:  # numpydoc ignore=PR01,RT01
     """See docstring in `array_api_extra._delegation.py`."""

src/array_api_extra/_lib/_funcs.py

@@ -166,8 +166,8 @@ def xp_assert_close(
         import numpy as np  # pylint: disable=import-outside-toplevel
 
         if is_pydata_sparse_namespace(xp):
-            actual = actual.to_dense()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
-            desired = desired.to_dense()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
+            actual = actual.todense()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
+            desired = desired.todense()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
 
         # JAX uses `np.testing`
         assert isinstance(rtol, float)

src/array_api_extra/_lib/_testing.py

@@ -1,4 +1,3 @@
-import contextlib
 import math
 import warnings
 from types import ModuleType
@@ -24,7 +23,7 @@
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
 from array_api_extra._lib._utils._compat import device as get_device
 from array_api_extra._lib._utils._helpers import eager_shape, ndindex
-from array_api_extra._lib._utils._typing import Array, Device
+from array_api_extra._lib._utils._typing import Device
 from array_api_extra.testing import lazy_xp_function
 
 # some xp backends are untyped
@@ -42,7 +41,6 @@
 lazy_xp_function(sinc, static_argnames="xp")
 
 
-@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no expand_dims")
 class TestAtLeastND:
     def test_0D(self, xp: ModuleType):
         x = xp.asarray(1.0)
@@ -69,7 +67,7 @@ def test_1D(self, xp: ModuleType):
         xp_assert_equal(y, xp.asarray([[0, 1]]))
 
         y = atleast_nd(x, ndim=5)
-        xp_assert_equal(y, xp.reshape(xp.arange(2), (1, 1, 1, 1, 2)))
+        xp_assert_equal(y, xp.asarray([[[[[0, 1]]]]]))
 
     def test_2D(self, xp: ModuleType):
         x = xp.asarray([[3.0]])
@@ -218,8 +216,10 @@ def test_xp(self, xp: ModuleType):
         )
 
 
+@pytest.mark.skip_xp_backend(
+    Backend.SPARSE, reason="read-only backend without .at support"
+)
 class TestCreateDiagonal:
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in zeros()")
     def test_1d_from_numpy(self, xp: ModuleType):
         # from np.diag tests
         vals = 100 * xp.arange(5, dtype=xp.float64)
@@ -235,7 +235,6 @@ def test_1d_from_numpy(self, xp: ModuleType):
         xp_assert_equal(create_diagonal(vals, offset=2), b)
         xp_assert_equal(create_diagonal(vals, offset=-2), c)
 
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in zeros()")
     @pytest.mark.parametrize("n", range(1, 10))
     @pytest.mark.parametrize("offset", range(1, 10))
     def test_1d_from_scipy(self, xp: ModuleType, n: int, offset: int):
@@ -251,7 +250,6 @@ def test_0d_raises(self, xp: ModuleType):
         with pytest.raises(ValueError, match="1-dimensional"):
             _ = create_diagonal(xp.asarray(1))
 
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in zeros()")
     @pytest.mark.parametrize(
         "shape",
         [
@@ -277,38 +275,24 @@ def test_nd(self, xp: ModuleType, shape: tuple[int, ...]):
         for i in ndindex(*eager_shape(c)):
             xp_assert_equal(c[i], b[i[:-1]] if i[-2] == i[-1] else zero)
 
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in zeros()")
     def test_device(self, xp: ModuleType, device: Device):
         x = xp.asarray([1, 2, 3], device=device)
         assert get_device(create_diagonal(x)) == device
 
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in zeros()")
     def test_xp(self, xp: ModuleType):
         x = xp.asarray([1, 2])
         y = create_diagonal(x, xp=xp)
         xp_assert_equal(y, xp.asarray([[1, 0], [0, 2]]))
 
 
 class TestExpandDims:
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no expand_dims")
-    @pytest.mark.xfail_xp_backend(Backend.DASK, reason="tuple index out of range")
-    @pytest.mark.xfail_xp_backend(Backend.TORCH, reason="tuple index out of range")
-    def test_functionality(self, xp: ModuleType):
-        def _squeeze_all(b: Array) -> Array:
-            """Mimics `np.squeeze(b)`. `xpx.squeeze`?"""
-            for axis in range(b.ndim):
-                with contextlib.suppress(ValueError):
-                    b = xp.squeeze(b, axis=axis)
-            return b
-
-        s = (2, 3, 4, 5)
-        a = xp.empty(s)
+    def test_single_axis(self, xp: ModuleType):
+        """Trivial case where xpx.expand_dims doesn't add anything to xp.expand_dims"""
+        a = xp.empty((2, 3, 4, 5))
         for axis in range(-5, 4):
             b = expand_dims(a, axis=axis)
-            assert b.shape[axis] == 1
-            assert _squeeze_all(b).shape == s
+            xp_assert_equal(b, xp.expand_dims(a, axis=axis))
 
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no expand_dims")
     def test_axis_tuple(self, xp: ModuleType):
         a = xp.empty((3, 3, 3))
         assert expand_dims(a, axis=(0, 1, 2)).shape == (1, 1, 1, 3, 3, 3)
@@ -317,8 +301,7 @@ def test_axis_tuple(self, xp: ModuleType):
         assert expand_dims(a, axis=(0, -3, -5)).shape == (1, 1, 3, 1, 3, 3)
 
     def test_axis_out_of_range(self, xp: ModuleType):
-        s = (2, 3, 4, 5)
-        a = xp.empty(s)
+        a = xp.empty((2, 3, 4, 5))
         with pytest.raises(IndexError, match="out of bounds"):
             _ = expand_dims(a, axis=-6)
         with pytest.raises(IndexError, match="out of bounds"):
@@ -341,12 +324,10 @@ def test_positive_negative_repeated(self, xp: ModuleType):
         with pytest.raises(ValueError, match="Duplicate dimensions"):
             _ = expand_dims(a, axis=(3, -3))
 
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no expand_dims")
     def test_device(self, xp: ModuleType, device: Device):
         x = xp.asarray([1, 2, 3], device=device)
         assert get_device(expand_dims(x, axis=0)) == device
 
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no expand_dims")
     def test_xp(self, xp: ModuleType):
         x = xp.asarray([1, 2, 3])
         y = expand_dims(x, axis=(0, 1, 2), xp=xp)
@@ -513,7 +494,6 @@ def test_xp(self, xp: ModuleType):
         xp_assert_equal(isclose(a, b, xp=xp), xp.asarray([True, False]))
 
 
-@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no expand_dims")
 class TestKron:
     def test_basic(self, xp: ModuleType):
         # Using 0-dimensional array
@@ -572,6 +552,7 @@ def test_kron_shape(
         k = kron(a, b)
         assert k.shape == expected_shape
 
+    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no isdtype")
     def test_python_scalar(self, xp: ModuleType):
         a = 1
         # Test no dtype promotion to xp.asarray(a); use b.dtype
@@ -614,25 +595,27 @@ def test_xp(self, xp: ModuleType):
         xp_assert_equal(nunique(a, xp=xp), xp.asarray(3))
 
 
-@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange, no device")
 class TestPad:
     def test_simple(self, xp: ModuleType):
-        a = xp.arange(1, 4)
+        a = xp.asarray([1, 2, 3])
         padded = pad(a, 2)
         xp_assert_equal(padded, xp.asarray([0, 0, 1, 2, 3, 0, 0]))
 
+    @pytest.mark.xfail_xp_backend(
+        Backend.SPARSE, reason="constant_values can only be equal to fill value"
+    )
     def test_fill_value(self, xp: ModuleType):
-        a = xp.arange(1, 4)
+        a = xp.asarray([1, 2, 3])
         padded = pad(a, 2, constant_values=42)
         xp_assert_equal(padded, xp.asarray([42, 42, 1, 2, 3, 42, 42]))
 
     def test_ndim(self, xp: ModuleType):
-        a = xp.reshape(xp.arange(2 * 3 * 4), (2, 3, 4))
+        a = xp.asarray(np.reshape(np.arange(2 * 3 * 4), (2, 3, 4)))
         padded = pad(a, 2)
         assert padded.shape == (6, 7, 8)
 
     def test_mode_not_implemented(self, xp: ModuleType):
-        a = xp.arange(3)
+        a = xp.asarray([1, 2, 3])
         with pytest.raises(NotImplementedError, match="Only `'constant'`"):
             _ = pad(a, 2, mode="edge")  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 
@@ -645,7 +628,7 @@ def test_xp(self, xp: ModuleType):
         xp_assert_equal(padded, xp.asarray(0))
 
     def test_tuple_width(self, xp: ModuleType):
-        a = xp.reshape(xp.arange(12), (3, 4))
+        a = xp.asarray(np.reshape(np.arange(12), (3, 4)))
         padded = pad(a, (1, 0))
         assert padded.shape == (4, 5)
 
@@ -656,7 +639,7 @@ def test_tuple_width(self, xp: ModuleType):
             _ = pad(a, [(1, 2, 3)])  # type: ignore[list-item]  # pyright: ignore[reportArgumentType]
 
     def test_sequence_of_tuples_width(self, xp: ModuleType):
-        a = xp.reshape(xp.arange(12), (3, 4))
+        a = xp.asarray(np.reshape(np.arange(12), (3, 4)))
 
         padded = pad(a, ((1, 0), (0, 2)))
         assert padded.shape == (4, 6)
@@ -678,7 +661,7 @@ def test_sequence_of_tuples_width(self, xp: ModuleType):
 )
 
 
-@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in asarray()")
+@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no argsort")
 class TestSetDiff1D:
     @pytest.mark.xfail_xp_backend(Backend.DASK, reason="NaN-shaped arrays")
     @pytest.mark.xfail_xp_backend(

tests/test_funcs.py

@@ -17,10 +17,8 @@
 lazy_xp_function(in1d, jax_jit=False, static_argnames=("assume_unique", "invert", "xp"))
 
 
+@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no unique_inverse")
 class TestIn1D:
-    @pytest.mark.xfail_xp_backend(
-        Backend.SPARSE, reason="no unique_inverse, no device kwarg in asarray()"
-    )
     # cover both code paths
     @pytest.mark.parametrize(
         "n",
@@ -42,19 +40,15 @@ def test_no_invert_assume_unique(self, xp: ModuleType, n: int):
         actual = in1d(x1, x2)
         xp_assert_equal(actual, expected)
 
-    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no device kwarg in asarray")
     def test_device(self, xp: ModuleType, device: Device):
         x1 = xp.asarray([3, 8, 20], device=device)
         x2 = xp.asarray([2, 3, 4], device=device)
         assert get_device(in1d(x1, x2)) == device
 
     @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp")
-    @pytest.mark.xfail_xp_backend(
-        Backend.SPARSE, reason="no arange, no device kwarg in asarray"
-    )
     def test_xp(self, xp: ModuleType):
         x1 = xp.asarray([1, 6])
-        x2 = xp.arange(5)
+        x2 = xp.asarray([0, 1, 2, 3, 4])
         expected = xp.asarray([True, False])
         actual = in1d(x1, x2, xp=xp)
         xp_assert_equal(actual, expected)
@@ -90,7 +84,7 @@ class TestAsArrays:
         ],
     )
     def test_array_vs_scalar(
-        self, dtype: str, b: int | float | complex, defined: bool, xp: ModuleType
+        self, dtype: str, b: complex, defined: bool, xp: ModuleType
     ):
         a = xp.asarray(1, dtype=getattr(xp, dtype))
 
@@ -158,7 +152,7 @@ def test_ndindex(shape: tuple[int, ...]):
     assert tuple(ndindex(*shape)) == tuple(np.ndindex(*shape))
 
 
-@pytest.mark.skip_xp_backend(Backend.SPARSE, reason="index by sparse array")
+@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="index by sparse array")
 def test_eager_shape(xp: ModuleType, library: Backend):
     a = xp.asarray([1, 2, 3])
     # Lazy arrays, like Dask, have an eager shape until you slice them with

tests/test_helpers.py

@@ -72,7 +72,7 @@ def test_assert_close_tolerance(xp: ModuleType):
 
 
 @param_assert_equal_close
-@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no bool indexing")
+@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="index by sparse array")
 def test_assert_close_equal_none_shape(xp: ModuleType, func: Callable[..., None]):  # type: ignore[no-any-explicit]
     """On dask and other lazy backends, test that a shape with NaN's or None's
     can be compared to a real shape.