Mlx

.github/workflows/mlx-tests.yml

@@ -0,0 +1,24 @@
+name: MLX Tests
+
+on: [push, pull_request]
+
+jobs:
+  mlx-tests:
+    runs-on: macos-14 # Apple Silicon runner
+    steps:
+      - uses: actions/checkout@v7.0.0
+
+      - name: Set up Python
+        uses: actions/setup-python@v6.2.0
+        with:
+          python-version: '3.13'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install pytest mlx numpy array-api-strict ndonnx
+
+      - name: Run MLX Compat Tests
+        run: |
+          export PYTHONPATH="${GITHUB_WORKSPACE}/src"
+          pytest tests/test_mlx.py tests/test_common.py tests/test_isdtype.py tests/test_copies_or_views.py -k mlx -v

meson.build

@@ -61,6 +61,15 @@ sources_raw = {
     'src/array_api_compat/torch/fft.py',
     'src/array_api_compat/torch/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/_typing.py',
+    'src/array_api_compat/mlx/fft.py',
+    'src/array_api_compat/mlx/linalg.py',
+  ],
 }
 
 sources = {}

pyproject.toml

@@ -55,6 +55,7 @@ dev = [
     "pytest",
     "torch",
     "sparse>=0.15.1",
+     "mlx; platform_system == 'Darwin' and platform_machine == 'arm64'"
 ]
 
 

src/array_api_compat/common/_helpers.py

@@ -116,7 +116,7 @@ def is_numpy_array(x: object) -> TypeIs[npt.NDArray[Any]]:
     # TODO: Should we reject ndarray subclasses?
     cls = cast(Hashable, type(x))
     return (
-        _issubclass_fast(cls, "numpy", "ndarray") 
+        _issubclass_fast(cls, "numpy", "ndarray")
         or _issubclass_fast(cls, "numpy", "generic")
     ) and not _is_jax_zero_gradient_array(x)
 
@@ -273,6 +273,29 @@ def is_pydata_sparse_array(x: object) -> TypeIs[sparse.SparseArray]:
     return _issubclass_fast(cls, "sparse", "SparseArray")
 
 
+def is_mlx_array(x: object) -> bool:
+    """
+    Return True if `x` is an MLX array.
+
+    This function does not import MLX 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_dask_array
+    is_jax_array
+    is_pydata_sparse_array
+    """
+    cls = cast(Hashable, type(x))
+    return _issubclass_fast(cls, "mlx.core", "array")
+
+
 def is_array_api_obj(x: object) -> TypeGuard[_ArrayApiObj]:
     """
     Return True if `x` is an array API compatible array object.
@@ -313,6 +336,7 @@ def _is_array_api_cls(cls: type) -> bool:
         # TODO: drop support for jax<0.4.32 which didn't have __array_namespace__
         or _issubclass_fast(cls, "jax", "Array")
         or _issubclass_fast(cls, "jax.core", "Tracer")  # see is_jax_array for limitations
+        or _issubclass_fast(cls, "mlx.core", "array")
     )
 
 
@@ -488,6 +512,27 @@ def is_array_api_strict_namespace(xp: Namespace) -> bool:
     return xp.__name__ == "array_api_strict"
 
 
+@lru_cache(100)
+def is_mlx_namespace(xp: Namespace) -> bool:
+    """
+    Returns True if `xp` is an MLX namespace.
+
+    This includes both ``mlx.core`` itself and the version wrapped by array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    """
+    return xp.__name__ in {"mlx.core", _compat_module_name() + ".mlx"}
+
+
 def _check_api_version(api_version: str | None) -> None:
     if api_version in _API_VERSIONS_OLD:
         warnings.warn(
@@ -516,7 +561,7 @@ def _cls_to_namespace(
     cls_ = cast(Hashable, cls)  # Make mypy happy
 
     if (
-        _issubclass_fast(cls_, "numpy", "ndarray") 
+        _issubclass_fast(cls_, "numpy", "ndarray")
         or _issubclass_fast(cls_, "numpy", "generic")
     ):
         if use_compat is True:
@@ -559,6 +604,14 @@ def _cls_to_namespace(
             import dask.array as xp  # type: ignore[no-redef]
         return xp, None
 
+    if _issubclass_fast(cls_, "mlx.core", "array"):
+        if _use_compat:
+            _check_api_version(api_version)
+            from .. import mlx as xp  # type: ignore[no-redef]
+        else:
+            import mlx.core as xp  # type: ignore[no-redef]
+        return xp, None
+
     # Backwards compatibility for jax<0.4.32
     if _issubclass_fast(cls_, "jax", "Array"):
         return _jax_namespace(api_version, use_compat), None
@@ -731,12 +784,45 @@ def __repr__(self) -> Literal["DASK_DEVICE"]:
 _DASK_DEVICE = _dask_device()
 
 
-# device() is not on numpy.ndarray or dask.array and to_device() is not on numpy.ndarray
-# or cupy.ndarray. They are not included in array objects of this library
-# because this library just reuses the respective ndarray classes without
-# wrapping or subclassing them. These helper functions can be used instead of
-# the wrapper functions for libraries that need to support both NumPy/CuPy and
-# other libraries that use devices.
+_mlx_device_map = {}
+
+
+def _mlx_cleanup(ref_id):
+    _mlx_device_map.pop(ref_id, None)
+
+
+def _get_mlx_device(x):
+    import mlx.core as mx
+    return _mlx_device_map.get(id(x), mx.default_device())
+
+
+def _set_mlx_device(x, device):
+    import weakref
+    ref_id = id(x)
+    _mlx_device_map[ref_id] = device
+    try:
+        # Register callback to pop ref_id when x is garbage collected.
+        # This prevents memory leaks without keeping a strong reference to x.
+        weakref.finalize(x, _mlx_cleanup, ref_id)
+    except TypeError:
+        # If x doesn't support weak references (though mlx arrays do)
+        pass
+
+
+def _normalize_mlx_device(device):
+    import mlx.core as mx
+    if isinstance(device, str):
+        if device == "cpu":
+            return mx.Device(mx.DeviceType.cpu)
+        elif device == "gpu":
+            return mx.Device(mx.DeviceType.gpu)
+        else:
+            raise ValueError(f"Unsupported device: {device!r}")
+    if not isinstance(device, mx.Device):
+        raise TypeError(f"Unsupported device type: {type(device)}")
+    return device
+
+
 def device(x: _ArrayApiObj, /) -> Device:
     """
     Hardware device the array data resides on.
@@ -802,6 +888,8 @@ def device(x: _ArrayApiObj, /) -> Device:
             return "cpu"
         # Return the device of the constituent array
         return device(inner)  # pyright: ignore
+    elif is_mlx_array(x):
+        return _get_mlx_device(x)
     return x.device  # type: ignore  # pyright: ignore
 
 
@@ -927,6 +1015,16 @@ def to_device(x: Array, device: Device, /, *, stream: int | Any | None = None) -
             if not hasattr(x, "to_device"):
                 return x
         return x.to_device(device, stream=stream)
+    elif is_mlx_array(x):
+        if stream is not None:
+            raise ValueError("The stream argument to to_device() is not supported for MLX")
+        dev = _normalize_mlx_device(device)
+        import mlx.core as mx
+        # MLX unified memory, so we just construct a new array reference
+        # and store its mapped device in the id->device registry (_mlx_device_map).
+        y = mx.array(x)
+        _set_mlx_device(y, dev)
+        return y
     elif is_pydata_sparse_array(x) and device == _device(x):
         # Perform trivial check to return the same array if
         # device is same instead of err-ing.
@@ -1079,6 +1177,8 @@ def is_lazy_array(x: object) -> TypeGuard[_ArrayApiObj]:
     "is_dask_namespace",
     "is_jax_array",
     "is_jax_namespace",
+    "is_mlx_array",
+    "is_mlx_namespace",
     "is_numpy_array",
     "is_numpy_namespace",
     "is_torch_array",

src/array_api_compat/mlx/__init__.py

@@ -0,0 +1,77 @@
+from typing import Final
+
+from .._internal import clone_module
+
+
+__all__ = clone_module("mlx.core", globals())
+
+# These imports may overwrite names from the import * above.
+from . import _aliases
+from ._aliases import *  # type: ignore[assignment,no-redef] # noqa: F403
+from ._info import __array_namespace_info__  # noqa: F401
+
+# Don't know why, but we have to do an absolute import to import linalg. If we
+# instead do
+#
+# from . import linalg
+#
+# It doesn't overwrite np.linalg from above. The import is generated
+# dynamically so that the library can be vendored.
+__import__(__spec__.parent + ".linalg")
+
+__import__(__spec__.parent + ".fft")
+
+from .linalg import matrix_transpose, vecdot  # type: ignore[no-redef]  # noqa: F401
+
+__array_api_version__: Final = "2025.12"
+
+__all__ = sorted(
+    set(__all__)
+    | set(_aliases.__all__)
+    | {"__array_api_version__", "__array_namespace_info__", "linalg", "fft"}
+)
+
+def __dir__() -> list[str]:
+    return __all__
+
+
+# Monkeypatch mlx.core.array to support boolean indexing __getitem__
+# MLX doesn't natively support boolean indexing __getitem__ (data-dependent shapes),
+# so we provide a NumPy-based fallback.
+import mlx.core as mx
+import numpy as np
+import builtins
+
+_old_getitem = mx.array.__getitem__
+
+def _is_bool_array(x):
+    return isinstance(x, mx.array) and x.dtype == mx.bool_
+
+def _has_bool_array(item):
+    if _is_bool_array(item):
+        return True
+    if isinstance(item, tuple):
+        return builtins.any(_has_bool_array(i) for i in item)
+    if isinstance(item, list):
+        return builtins.any(_has_bool_array(i) for i in item)
+    return False
+
+def _to_numpy_indices(item):
+    if _is_bool_array(item):
+        return np.array(item)
+    if isinstance(item, tuple):
+        return tuple(_to_numpy_indices(i) for i in item)
+    if isinstance(item, list):
+        return [_to_numpy_indices(i) for i in item]
+    return item
+
+def _new_getitem(self, item):
+    if _has_bool_array(item):
+        # Convert self and any boolean masks to numpy
+        self_np = np.array(self)
+        item_np = _to_numpy_indices(item)
+        res_np = self_np[item_np]
+        return mx.array(res_np)
+    return _old_getitem(self, item)
+
+mx.array.__getitem__ = _new_getitem