ENH: add mparray backend (#963)

* ENH: add mparray backend

[skip ci]

* Fixups

Author: mdhaber

pixi.lock

@@ -17,6 +17,7 @@ exclude-newer = "7d"
 
 [exclude-newer]
 pixi-build-python = "5d"  # for compatibility with pixi-build-api-version 7
+mparray = "0d"
 
 ### array-api-extra package definition ###
 
@@ -372,6 +373,7 @@ numpy = "=1.24.1"
 pytorch = ">=2.12.0"
 dask-core = ">=2026.7.1"  # No distributed, tornado, etc.
 sparse = ">=0.19.2"
+mparray = ">=0.2.2"
 
 [feature.backends.target.unix.dependencies]
 jax = ">=0.10.2"  # waiting for conda-forge/jaxlib-feedstock#326

pixi.toml

@@ -28,6 +28,7 @@ class Backend(Enum):  # numpydoc ignore=PR02
     ARRAY_API_STRICTEST = "array_api_strict:strictest"
     NUMPY = "numpy"
     NUMPY_READONLY = "numpy:readonly"
+    MPARRAY = "mparray"
     CUPY = "cupy"
     TORCH = "torch"
     TORCH_GPU = "torch:gpu"

src/array_api_extra/_lib/_backends.py

@@ -34,6 +34,7 @@
     "eager_shape",
     "in1d",
     "is_jax_jit_enabled",
+    "is_mparray_namespace",
     "is_python_scalar",
     "jax_autojit",
     "meta_namespace",
@@ -616,3 +617,8 @@ def is_jax_jit_enabled(xp: ArrayNamespace) -> bool:  # numpydoc ignore=PR01,RT01
         return bool(x)
     except jax.errors.TracerBoolConversionError:
         return True
+
+
+def is_mparray_namespace(xp: ArrayNamespace) -> bool:  # numpydoc ignore=PR01,RT01
+    """Return True if the argument is the MPArray namespace."""
+    return xp.__name__ == "mparray"

src/array_api_extra/_lib/_helpers.py

@@ -101,6 +101,7 @@ def nunique(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array:
     if (
         _compat.is_numpy_namespace(xp)
         or _compat.is_cupy_namespace(xp)
+        or _helpers.is_mparray_namespace(xp)
         or (
             _compat.is_torch_namespace(xp)
             and _helpers.capabilities(xp, x)["data-dependent shapes"]

src/array_api_extra/_set.py

@@ -695,6 +695,9 @@ def _as_numpy_array(  # numpydoc ignore=PR01,RT01
         cpu = typing.cast(Device, jax.devices("cpu")[0])
         array = _compat.to_device(array, cpu)
 
+    if _helpers.is_mparray_namespace(xp):
+        return np.asarray(array._data, dtype=array.dtype)  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
+
     if hasattr(array, "__dlpack__"):
         try:
             return np.from_dlpack(array)

src/array_api_extra/testing/_testing.py

@@ -246,7 +246,8 @@ def test_incompatible_dtype(
             with pytest.warns(FutureWarning, match="cannot safely cast"):
                 z = at_op(x, idx, op, 1.1, copy=copy)
 
-    elif library.like(Backend.DASK):
+    # MPArray mutation is currently a little too flexible; see mdhaber/mparray#21
+    elif library.like(Backend.DASK) or (library == Backend.MPARRAY):
         z = at_op(x, idx, op, 1.1, copy=copy)
 
     elif library.like(Backend.ARRAY_API_STRICT):

tests/main/test_at.py

@@ -662,6 +662,7 @@ def test_basic(self, xp: ArrayNamespace):
             atol=1e-11,
         )
 
+    @pytest.mark.skip_xp_backend(Backend.MPARRAY, reason="negative zero not supported")
     def test_real(self, xp: ArrayNamespace):
         x = xp.asarray([0.0, -0.0, 1.0, -1.0])
         expected = xp.asarray([0.0, xp.pi, 0.0, xp.pi], dtype=x.dtype)

tests/main/test_elementwise.py

@@ -27,6 +27,7 @@
                     Backend.TORCH_GPU, reason="device->host copy"
                 ),
                 pytest.mark.skip_xp_backend(Backend.SPARSE, reason="densification"),
+                pytest.mark.skip_xp_backend(Backend.MPARRAY, reason="precision loss"),
             ],
         ),
     ],

tests/main/test_lazy.py

@@ -144,6 +144,7 @@ def xp_searchsorted(
 
 @pytest.mark.skip_xp_backend(Backend.DASK, reason="no take_along_axis")
 @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="no searchsorted")
+@pytest.mark.skip_xp_backend(Backend.MPARRAY, reason="see mdhaber/mparray#20")
 class TestSearchsorted:
     def test_input_validation(self, xp: ArrayNamespace):
         message = "`side` must be either 'left' or 'right'."