From b3598dc14d22bf9b47862b3f3bb23eaf17401d38 Mon Sep 17 00:00:00 2001 From: Matt Haberland Date: Thu, 3 Sep 2026 17:17:52 -0700 Subject: [PATCH] ENH: Add mparray helper support Add mparray to backend detection, namespace discovery, device helpers, optional dependencies, CI, documentation, and shared tests. Treat it as a native CPU-only Array API backend. This commit was prepared with assistance from OpenAI Codex. --- .github/workflows/tests.yml | 6 +++ README.md | 3 +- docs/dev/tests.md | 5 +- docs/helper-functions.rst | 2 + docs/index.md | 14 ++--- docs/supported-array-libraries.md | 10 ++++ pyproject.toml | 4 +- src/array_api_compat/__init__.py | 2 +- src/array_api_compat/common/_helpers.py | 72 ++++++++++++++++++++++++- tests/_helpers.py | 2 +- tests/test_array_namespace.py | 4 +- tests/test_common.py | 8 ++- tests/test_no_dependencies.py | 2 +- 13 files changed, 116 insertions(+), 18 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index b2aa2348..56057cce 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -58,6 +58,12 @@ jobs: fi fi + # mparray requires Python >=3.11 and NumPy >=2. + if [[ "${{ matrix.python-version }}" =~ ^3\.1[1-4]$ && + "${{ matrix.numpy-version }}" =~ ^(latest|dev)$ ]]; then + python -m pip install mparray + fi + - name: Dump pip environment run: pip freeze diff --git a/README.md b/README.md index 3d80c708..4d6728c6 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,8 @@ This is a small wrapper around common array libraries that is compatible with the [Array API standard](https://data-apis.org/array-api/latest/). Currently, -NumPy, CuPy, PyTorch, Dask, JAX, ndonnx, `sparse` and dpnp are supported. If you want +NumPy, CuPy, PyTorch, Dask, JAX, mparray, ndonnx, `sparse` and dpnp are +supported. If you want support for other array libraries, or if you encounter any issues, please [open an issue](https://github.com/data-apis/array-api-compat/issues). diff --git a/docs/dev/tests.md b/docs/dev/tests.md index 18fb7cf5..d22c86ad 100644 --- a/docs/dev/tests.md +++ b/docs/dev/tests.md @@ -10,8 +10,9 @@ that is not strictly required by the standard. To run these tests, install the dependencies from the `dev` optional group (array-api-compat has [no hard runtime dependencies](no-dependencies)). -array-api-tests is run against all supported libraries are tested on CI -([except for JAX](jax-support) and [Sparse](sparse-support)). This is achieved +The array-api-tests suite is run against all supported libraries on CI +([except for JAX](jax-support), [mparray](mparray-support), and +[Sparse](sparse-support)). This is achieved by a [reusable GitHub Actions Workflow](https://github.com/data-apis/array-api-compat/blob/main/.github/workflows/array-api-tests.yml). Most libraries have tests that must be xfailed or skipped for various reasons. diff --git a/docs/helper-functions.rst b/docs/helper-functions.rst index d7ed6f6e..f04a4162 100644 --- a/docs/helper-functions.rst +++ b/docs/helper-functions.rst @@ -49,6 +49,7 @@ yet. .. autofunction:: is_torch_array .. autofunction:: is_dask_array .. autofunction:: is_jax_array +.. autofunction:: is_mparray_array .. autofunction:: is_pydata_sparse_array .. autofunction:: is_ndonnx_array .. autofunction:: is_dpnp_array @@ -59,6 +60,7 @@ yet. .. autofunction:: is_torch_namespace .. autofunction:: is_dask_namespace .. autofunction:: is_jax_namespace +.. autofunction:: is_mparray_namespace .. autofunction:: is_pydata_sparse_namespace .. autofunction:: is_ndonnx_namespace .. autofunction:: is_dpnp_namespace diff --git a/docs/index.md b/docs/index.md index baf2bee6..398ced0a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -2,7 +2,8 @@ This is a small wrapper around common array libraries that is compatible with the [Array API standard](https://data-apis.org/array-api/latest/). Currently, -NumPy, CuPy, PyTorch, Dask, JAX, ndonnx, Sparse and dpnp are supported. If you want +NumPy, CuPy, PyTorch, Dask, JAX, mparray, ndonnx, Sparse and dpnp are +supported. If you want support for other array libraries, or if you encounter any issues, please [open an issue](https://github.com/data-apis/array-api-compat/issues). @@ -61,12 +62,11 @@ import array_api_compat.dask as da ``` ```{note} -There are no `array_api_compat` submodules for JAX, sparse, or ndonnx. These -support for these libraries is contained in the libraries themselves (JAX -support is in the `jax.numpy` module in JAX v0.4.32 or newer, and in the -`jax.experimental.array_api` module for older JAX versions). The -array-api-compat support for these libraries consists of supporting them in -the [helper functions](helper-functions). +There are no `array_api_compat` submodules for JAX, mparray, sparse, or ndonnx. +These libraries provide their own Array API support. (JAX provides it through +`jax.numpy` in JAX v0.4.32+ or `jax.experimental.array_api` in older versions.) +array-api-compat supports these libraries through the +[helper functions](helper-functions). ``` Each will include all the functions from the normal NumPy/CuPy/PyTorch/dask.array diff --git a/docs/supported-array-libraries.md b/docs/supported-array-libraries.md index f5d89768..2d6fd227 100644 --- a/docs/supported-array-libraries.md +++ b/docs/supported-array-libraries.md @@ -100,6 +100,16 @@ Unlike the other libraries supported here, JAX array API support is contained entirely in the JAX library. The JAX array API support is tracked at https://github.com/google/jax/issues/18353. +(mparray-support)= +## [mparray](https://github.com/mdhaber/mparray) + +mparray implements the Array API directly, so there is no +`array_api_compat.mparray` wrapper. `array_namespace()` returns the native +`mparray` namespace, and compatibility support consists of library +detection and generic helper functions. mparray currently supports the +`2025.12` version of the standard and only the `"cpu"` device, so +`to_device()` follows NumPy's CPU-only behavior. + ## [Dask](https://www.dask.org/) If you're using dask with numpy, many of the same limitations that apply to numpy diff --git a/pyproject.toml b/pyproject.toml index 4ca2f429..7d52337d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ repository = "https://github.com/data-apis/array-api-compat/" cupy = ["cupy"] dask = ["dask>=2024.9.0"] jax = ["jax"] +mparray = ["mparray"] # Note: array-api-compat follows scikit-learn minimum dependencies, which support # much older versions of NumPy than what SPEC0 recommends. numpy = ["numpy>=1.22"] @@ -50,6 +51,7 @@ dev = [ "array-api-strict", "dask[array]>=2024.9.0", "jax[cpu]", + "mparray; python_version >= '3.11'", "ndonnx", "numpy>=1.22", "pytest", @@ -105,7 +107,7 @@ warn_unused_ignores = true warn_unreachable = true [[tool.mypy.overrides]] -module = ["cupy.*", "cupy_backends.*", "dask.*", "jax.*", "ndonnx.*", "sparse.*", "torch.*"] +module = ["cupy.*", "cupy_backends.*", "dask.*", "jax.*", "mparray.*", "ndonnx.*", "sparse.*", "torch.*"] ignore_missing_imports = true diff --git a/src/array_api_compat/__init__.py b/src/array_api_compat/__init__.py index eb3a948b..f7cef8d8 100644 --- a/src/array_api_compat/__init__.py +++ b/src/array_api_compat/__init__.py @@ -1,7 +1,7 @@ """ NumPy Array API compatibility library -This is a small wrapper around NumPy, CuPy, JAX, sparse and others that are +This is a small wrapper around NumPy, CuPy, JAX, mparray, sparse and others that are compatible with the Array API standard https://data-apis.org/array-api/latest/. See also NEP 47 https://numpy.org/neps/nep-0047-array-api-standard.html. diff --git a/src/array_api_compat/common/_helpers.py b/src/array_api_compat/common/_helpers.py index 76397cd4..3a7ae708 100644 --- a/src/array_api_compat/common/_helpers.py +++ b/src/array_api_compat/common/_helpers.py @@ -34,6 +34,7 @@ import dask.array as da import dpnp import jax + import mparray import ndonnx as ndx import numpy as np import numpy.typing as npt @@ -52,6 +53,7 @@ | dpnp.ndarray | dpnp.tensor.usm_ndarray | jax.Array + | mparray.MPArray | ndx.Array | sparse.SparseArray | torch.Tensor @@ -116,6 +118,7 @@ def is_numpy_array(x: object) -> TypeIs[npt.NDArray[Any]]: is_jax_array is_pydata_sparse_array is_dpnp_array + is_mparray_array """ # TODO: Should we reject ndarray subclasses? cls = cast(Hashable, type(x)) @@ -146,6 +149,7 @@ def is_cupy_array(x: object) -> bool: is_jax_array is_pydata_sparse_array is_dpnp_array + is_mparray_array """ cls = cast(Hashable, type(x)) return _issubclass_fast(cls, "cupy", "ndarray") @@ -169,6 +173,7 @@ def is_torch_array(x: object) -> TypeIs[torch.Tensor]: is_jax_array is_pydata_sparse_array is_dpnp_array + is_mparray_array """ cls = cast(Hashable, type(x)) return _issubclass_fast(cls, "torch", "Tensor") @@ -193,6 +198,7 @@ def is_ndonnx_array(x: object) -> TypeIs[ndx.Array]: is_jax_array is_pydata_sparse_array is_dpnp_array + is_mparray_array """ cls = cast(Hashable, type(x)) return _issubclass_fast(cls, "ndonnx", "Array") @@ -217,6 +223,7 @@ def is_dask_array(x: object) -> TypeIs[da.Array]: is_jax_array is_pydata_sparse_array is_dpnp_array + is_mparray_array """ cls = cast(Hashable, type(x)) return _issubclass_fast(cls, "dask.array", "Array") @@ -242,6 +249,7 @@ def is_jax_array(x: object) -> TypeIs[jax.Array]: is_dask_array is_pydata_sparse_array is_dpnp_array + is_mparray_array """ cls = cast(Hashable, type(x)) # We test for jax.core.Tracer here to identify jax arrays during jit tracing. From jax 0.8.2 on, @@ -257,6 +265,31 @@ def is_jax_array(x: object) -> TypeIs[jax.Array]: ) +def is_mparray_array(x: object) -> TypeIs[mparray.MPArray]: + """ + Return True if `x` is an mparray array. + + This function does not import mparray if it has not already been imported + and is therefore cheap to use. + + See Also + -------- + + array_namespace + is_array_api_obj + is_numpy_array + is_cupy_array + is_torch_array + is_ndonnx_array + is_dask_array + is_jax_array + is_pydata_sparse_array + is_dpnp_array + """ + cls = cast(Hashable, type(x)) + return _issubclass_fast(cls, "mparray", "MPArray") + + def is_pydata_sparse_array(x: object) -> TypeIs[sparse.SparseArray]: """ Return True if `x` is an array from the `sparse` package. @@ -277,6 +310,7 @@ def is_pydata_sparse_array(x: object) -> TypeIs[sparse.SparseArray]: is_dask_array is_jax_array is_dpnp_array + is_mparray_array """ # TODO: Account for other backends. cls = cast(Hashable, type(x)) @@ -303,6 +337,7 @@ def is_dpnp_array(x: object) -> bool: is_dask_array is_jax_array is_pydata_sparse_array + is_mparray_array """ cls = cast(Hashable, type(x)) return ( @@ -327,6 +362,7 @@ def is_array_api_obj(x: object) -> TypeGuard[_ArrayApiObj]: is_jax_array is_pydata_sparse_array is_dpnp_array + is_mparray_array """ try: # TODO: drop this check after np.matrix is gone @@ -379,6 +415,7 @@ def is_numpy_namespace(xp: Namespace) -> bool: is_jax_namespace is_pydata_sparse_namespace is_dpnp_namespace + is_mparray_namespace is_array_api_strict_namespace """ return xp.__name__ in {"numpy", _compat_module_name() + ".numpy"} @@ -402,6 +439,7 @@ def is_cupy_namespace(xp: Namespace) -> bool: is_jax_namespace is_pydata_sparse_namespace is_dpnp_namespace + is_mparray_namespace is_array_api_strict_namespace """ return xp.__name__ in {"cupy", _compat_module_name() + ".cupy"} @@ -425,6 +463,7 @@ def is_torch_namespace(xp: Namespace) -> bool: is_jax_namespace is_pydata_sparse_namespace is_dpnp_namespace + is_mparray_namespace is_array_api_strict_namespace """ return xp.__name__ in {"torch", _compat_module_name() + ".torch"} @@ -445,6 +484,7 @@ def is_ndonnx_namespace(xp: Namespace) -> bool: is_jax_namespace is_pydata_sparse_namespace is_dpnp_namespace + is_mparray_namespace is_array_api_strict_namespace """ return xp.__name__ == "ndonnx" @@ -468,6 +508,7 @@ def is_dask_namespace(xp: Namespace) -> bool: is_jax_namespace is_pydata_sparse_namespace is_dpnp_namespace + is_mparray_namespace is_array_api_strict_namespace """ return xp.__name__ in {"dask.array", _compat_module_name() + ".dask.array"} @@ -491,11 +532,33 @@ def is_jax_namespace(xp: Namespace) -> bool: is_dask_namespace is_pydata_sparse_namespace is_dpnp_namespace + is_mparray_namespace is_array_api_strict_namespace """ return xp.__name__ in {"jax.numpy", "jax.experimental.array_api"} +def is_mparray_namespace(xp: Namespace) -> bool: + """ + Returns True if `xp` is the mparray namespace. + + See Also + -------- + + array_namespace + is_numpy_namespace + is_cupy_namespace + is_torch_namespace + is_ndonnx_namespace + is_dask_namespace + is_jax_namespace + is_pydata_sparse_namespace + is_dpnp_namespace + is_array_api_strict_namespace + """ + return xp.__name__ == "mparray" + + def is_pydata_sparse_namespace(xp: Namespace) -> bool: """ Returns True if `xp` is a pydata/sparse namespace. @@ -511,6 +574,7 @@ def is_pydata_sparse_namespace(xp: Namespace) -> bool: is_dask_namespace is_jax_namespace is_dpnp_namespace + is_mparray_namespace is_array_api_strict_namespace """ return xp.__name__ == "sparse" @@ -533,6 +597,7 @@ def is_dpnp_namespace(xp: Namespace) -> bool: is_dask_namespace is_jax_namespace is_pydata_sparse_namespace + is_mparray_namespace is_array_api_strict_namespace """ return xp.__name__ in {"dpnp", "dpnp.tensor"} @@ -554,6 +619,7 @@ def is_array_api_strict_namespace(xp: Namespace) -> bool: is_jax_namespace is_pydata_sparse_namespace is_dpnp_namespace + is_mparray_namespace """ return xp.__name__ == "array_api_strict" @@ -714,6 +780,7 @@ def your_function(x, y): is_torch_array is_dask_array is_jax_array + is_mparray_array is_pydata_sparse_array is_dpnp_array @@ -971,7 +1038,7 @@ def to_device(x: Array, device: Device, /, *, stream: int | Any | None = None) - device : Hardware device the array data resides on. """ - if is_numpy_array(x): + if is_numpy_array(x) or is_mparray_array(x): if stream is not None: raise ValueError("The stream argument to to_device() is not supported") if device == "cpu": @@ -1070,6 +1137,7 @@ def _is_lazy_cls(cls: type) -> bool | None: or _issubclass_fast(cls, "cupy", "ndarray") or _issubclass_fast(cls, "torch", "Tensor") or _issubclass_fast(cls, "sparse", "SparseArray") + or _issubclass_fast(cls, "mparray", "MPArray") ): return False if ( @@ -1152,6 +1220,8 @@ def is_lazy_array(x: object) -> TypeGuard[_ArrayApiObj]: "is_dpnp_namespace", "is_jax_array", "is_jax_namespace", + "is_mparray_array", + "is_mparray_namespace", "is_numpy_array", "is_numpy_namespace", "is_torch_array", diff --git a/tests/_helpers.py b/tests/_helpers.py index b12ff720..4352cf3b 100644 --- a/tests/_helpers.py +++ b/tests/_helpers.py @@ -4,7 +4,7 @@ wrapped_libraries = ["numpy", "cupy", "torch", "dask.array"] all_libraries = wrapped_libraries + [ - "array_api_strict", "jax.numpy", "ndonnx", "sparse", "dpnp" + "array_api_strict", "jax.numpy", "ndonnx", "sparse", "dpnp", "mparray" ] def import_(library, wrapper=False): diff --git a/tests/test_array_namespace.py b/tests/test_array_namespace.py index e56b8eb4..0b1121dc 100644 --- a/tests/test_array_namespace.py +++ b/tests/test_array_namespace.py @@ -26,7 +26,9 @@ def test_array_namespace(request, library, api_version, use_compat): if (library == "sparse" and api_version in ("2023.12", "2024.12")) or ( library == "jax.numpy" and api_version in ("2021.12", "2022.12", "2023.12", "2024.12")) or ( - library == "dpnp" and api_version in ("2021.12", "2022.12", "2023.12") + library == "dpnp" and api_version in ("2021.12", "2022.12", "2023.12")) or ( + library == "mparray" and + api_version in ("2021.12", "2022.12", "2023.12", "2024.12") ): xfail(request, "Unsupported API version") diff --git a/tests/test_common.py b/tests/test_common.py index ff7c3d3f..9526fec6 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -8,11 +8,11 @@ from array_api_compat import ( # noqa: F401 is_numpy_array, is_cupy_array, is_torch_array, is_dask_array, is_jax_array, is_pydata_sparse_array, - is_ndonnx_array, is_dpnp_array, + is_ndonnx_array, is_dpnp_array, is_mparray_array, is_numpy_namespace, is_cupy_namespace, is_torch_namespace, is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace, is_array_api_strict_namespace, is_ndonnx_namespace, - is_dpnp_namespace, + is_dpnp_namespace, is_mparray_namespace, ) from array_api_compat import ( @@ -31,6 +31,7 @@ 'sparse': 'is_pydata_sparse_array', 'ndonnx': 'is_ndonnx_array', 'dpnp': 'is_dpnp_array', + 'mparray': 'is_mparray_array', } is_namespace_functions = { @@ -43,6 +44,7 @@ 'array_api_strict': 'is_array_api_strict_namespace', 'ndonnx': 'is_ndonnx_namespace', 'dpnp': 'is_dpnp_namespace', + 'mparray': 'is_mparray_namespace', } @@ -259,6 +261,8 @@ def test_asarray_cross_library(source_library, target_library, request): elif source_library == "jax.numpy" and target_library == "torch": xfail(request, reason="casts int to float") + elif source_library == "mparray" and target_library != "mparray": + pytest.skip(reason="mparray does not support implicit conversion") elif source_library == "cupy" and target_library != "cupy": # cupy explicitly disallows implicit conversions to CPU pytest.skip(reason="cupy does not support implicit conversion to CPU") diff --git a/tests/test_no_dependencies.py b/tests/test_no_dependencies.py index 624f8971..031b6240 100644 --- a/tests/test_no_dependencies.py +++ b/tests/test_no_dependencies.py @@ -50,7 +50,7 @@ def _test_dependency(mod): # Y (except most array libraries actually do themselves depend on numpy). @pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array", - "jax.numpy", "sparse", "array_api_strict"]) + "jax.numpy", "mparray", "sparse", "array_api_strict"]) def test_numpy_dependency(library): # This import is here because it imports numpy from ._helpers import import_