ENH: Adding MLX

.github/workflows/ci.yml

@@ -59,7 +59,8 @@ jobs:
   checks:
     name: ${{ matrix.environment }} (${{ matrix.platform }})
     runs-on: >-
-      ${{ matrix.platform == 'macos' && 'macos-latest'
+      ${{ matrix.platform == 'macos-mlx' && 'macos-15'
+       || matrix.platform == 'macos' && 'macos-latest'
        || matrix.platform == 'windows' && 'windows-latest'
        || matrix.platform == 'linux-arm' && 'ubuntu-24.04-arm'
        || matrix.x64_runner || 'ubuntu-slim' }}
@@ -90,6 +91,8 @@ jobs:
               no-coverage: true,
               x64_runner: ubuntu-latest,
             }
+          - environment: tests-mlx
+            platform: macos-mlx
           - {environment: tests-run-deps, task: tests-run-deps-cov}
         exclude:
           - {environment: tests-numpy1, platform: windows} # data-apis/array-api-extra#901

pixi.toml

@@ -7,6 +7,7 @@ platforms = [
   "linux-aarch64",
   "osx-64",
   "osx-arm64",
+  { name = "osx-arm64-macos14", platform = "osx-arm64", macos = "14.5" },
   "win-64",
   { platform = "linux-64", cuda = "12.9" },
   { platform = "win-64", cuda = "12.9" },
@@ -69,6 +70,7 @@ tests-backends = {
   solve-group = "backends",
 }
 tests-backends-py311.features = ["py311", "tests", "backends"]
+tests-mlx = { features = ["py314", "tests", "mlx"], solve-group = "mlx" }
 
 # CUDA not available on free github actions and on some developers' PCs
 dev-cuda = {
@@ -378,6 +380,12 @@ mparray = ">=0.2.2"
 [feature.backends.target.unix.dependencies]
 jax = ">=0.10.2"  # waiting for conda-forge/jaxlib-feedstock#326
 
+[feature.mlx]
+platforms = ["osx-arm64-macos14"]
+
+[feature.mlx.target.osx-arm64.pypi-dependencies]
+mlx = ">=0.32.2"
+
 # Backends that require a GPU host and a CUDA driver.
 [feature.cuda-backends]
 platforms = ["linux-64-cuda-12-9", "win-64-cuda-12-9"]

pyproject.toml

@@ -80,7 +80,7 @@ disable_error_code = ["no-any-return"]
 
 [[tool.mypy.overrides]]
 # slow or unavailable on Windows; do not add to the lint env
-module = ["cupy.*", "jax.*", "sparse.*", "torch.*"]
+module = ["cupy.*", "jax.*", "mlx.*", "sparse.*", "torch.*"]
 ignore_missing_imports = true
 
 [[tool.mypy.overrides]]

src/array_api_extra/_lib/_backends.py

@@ -30,6 +30,7 @@ class Backend(Enum):  # numpydoc ignore=PR02
     NUMPY_READONLY = "numpy:readonly"
     MPARRAY = "mparray"
     CUPY = "cupy"
+    MLX = "mlx.core"
     TORCH = "torch"
     TORCH_GPU = "torch:gpu"
     DASK = "dask.array"

src/array_api_extra/_lib/_compat.py

@@ -2,11 +2,28 @@
 # Allow packages that vendor both `array-api-extra` and
 # `array-api-compat` to override the import location
 
-# pylint: disable=duplicate-code
+from collections.abc import Callable
+from functools import cache
+from types import ModuleType
+from typing import TYPE_CHECKING, ClassVar, Protocol, cast
+
+if TYPE_CHECKING:  # pragma: no cover
+    from typing_extensions import override
+else:
+
+    def override(func):
+        return func
+
+
+# pylint: disable=duplicate-code,redefined-outer-name
 try:
     from ..._array_api_compat_vendor import (
-        array_namespace,
-        device,
+        array_namespace as _array_namespace,
+    )
+    from ..._array_api_compat_vendor import (
+        device as _device,
+    )
+    from ..._array_api_compat_vendor import (
         is_array_api_obj,
         is_array_api_strict_namespace,
         is_cupy_array,
@@ -24,12 +41,18 @@
         is_torch_namespace,
         is_writeable_array,
         size,
-        to_device,
+    )
+    from ..._array_api_compat_vendor import (
+        to_device as _to_device,
     )
 except ImportError:
     from array_api_compat import (
-        array_namespace,
-        device,
+        array_namespace as _array_namespace,
+    )
+    from array_api_compat import (
+        device as _device,
+    )
+    from array_api_compat import (
         is_array_api_obj,
         is_array_api_strict_namespace,
         is_cupy_array,
@@ -47,8 +70,149 @@
         is_torch_namespace,
         is_writeable_array,
         size,
-        to_device,
     )
+    from array_api_compat import (
+        to_device as _to_device,
+    )
+
+
+class _NamespaceInfo(Protocol):  # numpydoc ignore=GL08
+    def default_dtypes(self) -> object: ...  # numpydoc ignore=GL08
+
+
+class _MLXNamespaceInfo:  # numpydoc ignore=GL08
+    _info: _NamespaceInfo
+
+    def __init__(self, info: _NamespaceInfo) -> None:  # numpydoc ignore=GL08
+        self._info = info
+
+    def __getattr__(self, name: str) -> object:  # numpydoc ignore=GL08
+        return getattr(self._info, name)
+
+    def default_dtypes(
+        self, *, device: object = None
+    ) -> object:  # numpydoc ignore=GL08
+        _ = device
+        return self._info.default_dtypes()
+
+
+class _MLXNamespace(ModuleType):  # numpydoc ignore=GL08
+    _device_functions: ClassVar[set[str]] = {
+        "arange",
+        "asarray",
+        "empty",
+        "eye",
+        "full",
+        "linspace",
+        "ones",
+        "zeros",
+    }
+
+    _namespace: ModuleType
+
+    def __init__(self, namespace: ModuleType) -> None:  # numpydoc ignore=GL08
+        super().__init__(namespace.__name__)
+        self._namespace = namespace
+
+    @override
+    def __getattr__(self, name: str) -> object:  # numpydoc ignore=GL08
+        if name == "__array_api_version__":
+            return "2024.12"
+        if name == "bool":
+            return self._namespace.bool_
+        if name == "__array_namespace_info__":
+            get_info = cast(
+                Callable[[], _NamespaceInfo], getattr(self._namespace, name)
+            )
+            return lambda: _MLXNamespaceInfo(get_info())
+        if name == "signbit":
+            return lambda x: x < 0
+        function = getattr(self._namespace, name)
+        if name not in {"argsort", "astype", "result_type", "sort"} and (
+            name not in self._device_functions
+        ):
+            return function
+
+        def compatible_call(  # numpydoc ignore=GL08
+            *args: object,
+            device: object = None,
+            copy: bool | None = None,
+            **kwargs: object,
+        ) -> object:
+            if name == "full" and "fill_value" in kwargs:
+                args = (*args, kwargs.pop("fill_value"))
+            if name == "asarray" and copy is not None:
+                kwargs["copy"] = copy
+            if name == "asarray" and "dtype" not in kwargs:
+                input_array = args[0] if args else kwargs.get("a")
+                input_dtype = getattr(input_array, "dtype", None)
+                dtype_name = getattr(input_dtype, "name", None)
+                if dtype_name is not None and hasattr(self._namespace, dtype_name):
+                    kwargs["dtype"] = getattr(self._namespace, dtype_name)
+            if name == "astype":
+                _ = copy
+                return function(*args, **kwargs)
+            if name in {"argsort", "sort"}:
+                _ = kwargs.pop("stable", None)
+            if name == "result_type":
+                import mlx.core as mx
+
+                converted_args: list[object] = []
+                for value in args:
+                    if isinstance(value, int | float | complex | bool):
+                        converted_args.append(cast(object, mx.asarray(value)))
+                    else:
+                        converted_args.append(value)
+                args = tuple(converted_args)
+            if device is None:
+                return function(*args, **kwargs)
+            import mlx.core as mx
+
+            with mx.stream(device):
+                return function(*args, **kwargs)
+
+        return compatible_call
+
+
+@cache
+def _wrap_mlx_namespace(namespace: ModuleType) -> ModuleType:  # numpydoc ignore=GL08
+    return _MLXNamespace(namespace)
+
+
+def array_namespace(  # numpydoc ignore=GL08
+    *xs: object, api_version: str | None = None, use_compat: bool | None = None
+) -> ModuleType:
+    namespace = cast(
+        ModuleType,
+        _array_namespace(*xs, api_version=api_version, use_compat=use_compat),
+    )
+    if namespace.__name__ == "mlx.core":
+        return _wrap_mlx_namespace(namespace)
+    return namespace
+
+
+def device(x: object, /) -> object:  # numpydoc ignore=GL08
+    if type(x).__module__.startswith("mlx."):
+        import mlx.core as mx
+
+        return mx.default_device()
+    return _device(x)
+
+
+def to_device(  # numpydoc ignore=GL08
+    x: object, device: object, /, *, stream: object = None
+) -> object:
+    if type(x).__module__.startswith("mlx."):
+        import mlx.core as mx
+
+        if device == "cpu":
+            device = mx.cpu
+        elif device == "gpu":
+            device = mx.gpu
+        with mx.stream(device):
+            return mx.array(x)
+    return _to_device(x, device, stream=stream)
+
 
 __all__ = [
     "array_namespace",

tests/main/conftest.py

@@ -148,6 +148,14 @@ def xp(
         _setup_jax(library)
     elif library.like(Backend.TORCH):
         _setup_torch(library)
+    elif library is Backend.MLX:
+        _setup_mlx()
+
+        import mlx.core as mx
+
+        with mx.stream(mx.cpu), patch_lazy_xp_functions(request, xp=xp):
+            yield xp
+        return
 
     # On Dask and JAX, monkey-patch all functions tagged by `lazy_xp_function`
     # in the global scope of the module containing the test function.
@@ -189,6 +197,12 @@ def _setup_torch(library: Backend) -> None:
         torch.set_default_device("cpu")
 
 
+def _setup_mlx() -> None:
+    import mlx.core as mx
+
+    mx.set_default_device(mx.cpu)
+
+
 # Can select the test with `pytest -k dask`
 @pytest.fixture(params=[Backend.DASK.pytest_param()])
 def da(
@@ -244,6 +258,6 @@ def device(
 @pytest.fixture
 def infinity(library: Backend) -> float:
     """Retrieve the positive infinity value for the given backend."""
-    if library in (Backend.TORCH, Backend.TORCH_GPU):
+    if library in (Backend.MLX, Backend.TORCH, Backend.TORCH_GPU):
         return 3.4028235e38
     return 1.7976931348623157e308

tests/main/test_at.py

@@ -121,6 +121,7 @@ def assert_copy(
 )
 def test_update_ops(
     xp: ArrayNamespace,
+    library: Backend,
     copy: bool | None,
     op: _AtOp,
     y: float,
@@ -129,6 +130,8 @@ def test_update_ops(
     x_ndim: int,
     y_ndim: int,
 ):
+    if bool_mask and library is Backend.MLX:
+        pytest.skip("MLX does not support boolean indexing")
     if x_ndim == 1:
         x = xp.asarray([10.0, 20.0, 30.0])
         idx = xp.asarray([False, True, True]) if bool_mask else slice(1, None)
@@ -153,6 +156,7 @@ def test_update_ops(
 
 
 @pytest.mark.parametrize("op", list(_AtOp))
+@pytest.mark.skip_xp_backend(Backend.MLX, reason="boolean indexing is unsupported")
 def test_copy_default(xp: ArrayNamespace, library: Backend, op: _AtOp):
     """
     Test that the default copy behaviour is False for writeable arrays
@@ -235,6 +239,8 @@ def test_incompatible_dtype(
     UFuncTypeError: Cannot cast ufunc 'divide' output from dtype('float64')
     to dtype('int64') with casting rule 'same_kind'
     """
+    if bool_mask and library is Backend.MLX:
+        pytest.skip("MLX does not support boolean indexing")
     x = xp.asarray([2, 4])
     idx = xp.asarray([True, False]) if bool_mask else slice(None)
     z = None
@@ -254,6 +260,13 @@ def test_incompatible_dtype(
         with pytest.raises(Exception, match=r"cast|promote|dtype"):
             _ = at_op(x, idx, op, 1.1, copy=copy)
 
+    elif library is Backend.MLX:
+        if op is _AtOp.DIVIDE:
+            with pytest.raises(Exception, match="cast"):
+                _ = at_op(x, idx, op, 1.1, copy=copy)
+        else:
+            z = at_op(x, idx, op, 1.1, copy=copy)
+
     elif op in (_AtOp.SET, _AtOp.MIN, _AtOp.MAX):
         # There is no __i<op>__ version of min/max.
         # libraries other than array-api-strict are happy with
@@ -277,7 +290,9 @@ def test_bool_mask_nd(xp: ArrayNamespace):
 
 
 @pytest.mark.parametrize("bool_mask", [False, True])
-def test_no_inf_warnings(xp: ArrayNamespace, bool_mask: bool):
+def test_no_inf_warnings(xp: ArrayNamespace, library: Backend, bool_mask: bool):
+    if bool_mask and library is Backend.MLX:
+        pytest.skip("MLX does not support boolean indexing")
     x = xp.asarray([math.inf, 1.0, 2.0])
     idx = ~xp.isinf(x) if bool_mask else slice(1, None)
     # inf - inf -> nan with a warning

tests/main/test_creation.py

@@ -49,6 +49,9 @@ def test_2d(self, xp: ArrayNamespace):
     @pytest.mark.skip_xp_backend(
         Backend.ARRAY_API_STRICTEST, reason="backend doesn't support Boolean indexing"
     )
+    @pytest.mark.skip_xp_backend(
+        Backend.MLX, reason="backend doesn't support Boolean indexing"
+    )
     def test_abstract_size(self, xp: ArrayNamespace):
         x = xp.arange(5)
         x = x[x > 2]