Add an MLX Array API compatibility namespace

.github/workflows/array-api-tests-mlx.yml

@@ -0,0 +1,70 @@
+name: MLX Array API Tests
+
+on:
+  push:
+    branches:
+      - agent/mlx-compat-complete
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+
+concurrency:
+  group: mlx-array-api-${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  tests:
+    runs-on: macos-14
+    timeout-minutes: 60
+
+    steps:
+      - name: Checkout array-api-compat
+        uses: actions/checkout@v4
+
+      - name: Checkout array-api-tests
+        uses: actions/checkout@v4
+        with:
+          repository: data-apis/array-api-tests
+          submodules: true
+          path: array-api-tests
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install mlx pytest
+          python -m pip install -r array-api-tests/requirements.txt
+          python -m pip install .
+
+      - name: Run focused MLX wrapper tests
+        run: python -m pytest tests/test_mlx.py -q
+
+      - name: Run Array API conformance suite
+        env:
+          ARRAY_API_TESTS_MODULE: array_api_compat.mlx
+          ARRAY_API_TESTS_VERSION: "2025.12"
+        run: |
+          python - <<'PY'
+          import mlx.core as mx
+          import pytest
+
+          mx.set_default_device(mx.cpu)
+          raise SystemExit(
+              pytest.main(
+                  [
+                      "array-api-tests/array_api_tests",
+                      "--max-examples=100",
+                      "--hypothesis-disable-deadline",
+                      "--xfails-file=mlx-xfails.txt",
+                      "--skips-file=mlx-skips.txt",
+                      "-q",
+                      "-rxXfE",
+                  ]
+              )
+          )
+          PY

meson.build

@@ -20,6 +20,7 @@ sources_raw = {
     'src/array_api_compat/common/_fft.py',
     'src/array_api_compat/common/_helpers.py',
     'src/array_api_compat/common/_linalg.py',
+    'src/array_api_compat/common/_mlx_helpers.py',
     'src/array_api_compat/common/_typing.py',
   ],
 
@@ -44,6 +45,16 @@ sources_raw = {
     'src/array_api_compat/dask/array/linalg.py',
   ],
 
+  'array_api_compat/mlx': [
+    'src/array_api_compat/mlx/__init__.py',
+    'src/array_api_compat/mlx/_aliases.py',
+    'src/array_api_compat/mlx/_info.py',
+    'src/array_api_compat/mlx/_overrides.py',
+    'src/array_api_compat/mlx/_typing.py',
+    'src/array_api_compat/mlx/fft.py',
+    'src/array_api_compat/mlx/linalg.py',
+  ],
+
   'array_api_compat/numpy': [
     'src/array_api_compat/numpy/__init__.py',
     'src/array_api_compat/numpy/_aliases.py',

mlx-skips.txt

@@ -0,0 +1 @@
+# MLX tests are xfailed rather than skipped so unexpected passes remain visible.

mlx-xfails.txt

@@ -0,0 +1,18 @@
+# MLX arrays intentionally do not implement data-dependent output shapes.
+array_api_tests/test_array_object.py::test_getitem
+array_api_tests/test_searching_functions.py::test_nonzero
+array_api_tests/test_set_functions.py::test_unique_all
+array_api_tests/test_set_functions.py::test_unique_counts
+array_api_tests/test_set_functions.py::test_unique_inverse
+array_api_tests/test_set_functions.py::test_unique_values
+
+# array-api-compat does not wrap or monkeypatch mlx.core.array.
+array_api_tests/test_has_names.py::test_has_names[array_method-__index__]
+array_api_tests/test_has_names.py::test_has_names[array_method-to_device]
+array_api_tests/test_has_names.py::test_has_names[array_attribute-device]
+array_api_tests/test_has_names.py::test_has_names[array_attribute-mT]
+array_api_tests/test_signatures.py::test_array_method_signature[__index__]
+array_api_tests/test_signatures.py::test_array_method_signature[to_device]
+
+# MLX does not expose non-scalar, data-dependent repeat shapes.
+array_api_tests/test_manipulation_functions.py::test_repeat

pyproject.toml

@@ -32,6 +32,7 @@ repository = "https://github.com/data-apis/array-api-compat/"
 cupy = ["cupy"]
 dask = ["dask>=2024.9.0"]
 jax = ["jax"]
+mlx = ["mlx; sys_platform == 'darwin'"]
 # 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]",
+    "mlx; sys_platform == 'darwin'",
     "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.*", "mlx.*", "ndonnx.*", "sparse.*", "torch.*"]
 ignore_missing_imports = true
 
 

src/array_api_compat/common/__init__.py

@@ -1 +1,11 @@
 from ._helpers import *  # noqa: F403
+from ._mlx_helpers import (  # noqa: F401
+    array_namespace,
+    device,
+    get_namespace,
+    is_array_api_obj,
+    is_lazy_array,
+    is_mlx_array,
+    is_mlx_namespace,
+    to_device,
+)

src/array_api_compat/common/_mlx_helpers.py

@@ -0,0 +1,120 @@
+"""Lazy MLX integration for the public helper functions.
+
+This module intentionally does not import MLX merely because
+``array_api_compat`` is imported. MLX is imported only after an actual
+``mlx.core.array`` or MLX namespace has been supplied.
+"""
+
+from __future__ import annotations
+
+import sys
+from types import ModuleType
+from typing import Any
+
+from . import _helpers as _base
+
+_SCALAR_TYPES = (bool, int, float, complex, type(None))
+
+
+def is_mlx_array(x: object) -> bool:
+    """Return whether *x* is an MLX array without importing MLX."""
+    module = sys.modules.get("mlx.core")
+    if module is None:
+        return False
+    array_type = getattr(module, "array", None)
+    return array_type is not None and isinstance(x, array_type)
+
+
+def is_mlx_namespace(xp: ModuleType) -> bool:
+    """Return whether *xp* is MLX or the array-api-compat MLX wrapper."""
+    return xp.__name__ in {"mlx.core", _base._compat_module_name() + ".mlx"}
+
+
+def array_namespace(
+    *xs: Any,
+    api_version: str | None = None,
+    use_compat: bool | None = None,
+) -> ModuleType:
+    """Return the Array API namespace, including the MLX compat wrapper."""
+    mlx_inputs = [x for x in xs if is_mlx_array(x)]
+    if not mlx_inputs:
+        return _base.array_namespace(
+            *xs,
+            api_version=api_version,
+            use_compat=use_compat,
+        )
+
+    for x in xs:
+        if isinstance(x, _SCALAR_TYPES):
+            continue
+        if not is_mlx_array(x):
+            raise TypeError("Multiple namespaces for array inputs: MLX and another backend")
+
+    _base._check_api_version(api_version)
+    if use_compat is False:
+        import mlx.core as mx
+
+        return mx
+
+    from .. import mlx as mlx_compat
+
+    return mlx_compat
+
+
+get_namespace = array_namespace
+
+
+def is_array_api_obj(x: object) -> bool:
+    return is_mlx_array(x) or _base.is_array_api_obj(x)
+
+
+def device(x: Any, /) -> Any:
+    if not is_mlx_array(x):
+        return _base.device(x)
+
+    import mlx.core as mx
+
+    # MLX arrays use unified memory and do not carry per-array residency.
+    # The execution default is the only meaningful device value MLX exposes.
+    return mx.default_device()
+
+
+def to_device(
+    x: Any,
+    device: Any,
+    /,
+    *,
+    stream: int | Any | None = None,
+) -> Any:
+    if not is_mlx_array(x):
+        return _base.to_device(x, device, stream=stream)
+    if stream is not None:
+        raise NotImplementedError("MLX does not expose Array API stream handles")
+
+    import mlx.core as mx
+
+    if not isinstance(device, mx.Device):
+        raise TypeError(f"expected an mlx.core.Device, got {type(device).__name__}")
+    return mx.full_like(x, x, dtype=x.dtype, stream=device)
+
+
+def is_lazy_array(x: object) -> bool:
+    if is_mlx_array(x):
+        return True
+    return _base.is_lazy_array(x)
+
+
+__all__ = [
+    "array_namespace",
+    "device",
+    "get_namespace",
+    "is_array_api_obj",
+    "is_lazy_array",
+    "is_mlx_array",
+    "is_mlx_namespace",
+    "to_device",
+]
+
+
+def __dir__() -> list[str]:
+    return __all__

src/array_api_compat/mlx/__init__.py

@@ -0,0 +1,33 @@
+from typing import Final
+
+from .._internal import clone_module
+
+__all__ = clone_module("mlx.core", globals())
+
+from . import _aliases, _overrides
+from ._aliases import *  # type: ignore[assignment,no-redef] # noqa: F403
+from ._overrides import *  # type: ignore[assignment,no-redef] # noqa: F403
+from ._info import __array_namespace_info__
+
+# Import the compatibility submodules explicitly so they replace the native
+# ``mlx.core.fft`` and ``mlx.core.linalg`` objects cloned above.
+fft = __import__(__spec__.parent + ".fft", fromlist=["fft"])
+linalg = __import__(__spec__.parent + ".linalg", fromlist=["linalg"])
+
+__array_api_version__: Final = "2025.12"
+
+__all__ = sorted(
+    set(__all__)
+    | set(_aliases.__all__)
+    | set(_overrides.__all__)
+    | {
+        "__array_api_version__",
+        "__array_namespace_info__",
+        "fft",
+        "linalg",
+    }
+)
+
+
+def __dir__() -> list[str]:
+    return __all__