Mlx

[aaishwarymishra]

MLX Support Overview

MLX is built from the ground up for Apple Silicon using a Unified Memory model—the CPU and GPU share the exact same physical memory. This presents
unique design challenges when wrapping it for the Array API standard (which assumes distinct hardware platforms like CPU vs. GPU VRAM).

In general, array-api-compat wraps existing libraries without subclassing or wrapping their array objects, keeping it lightweight. However, because
MLX is missing a few core features required by the standard (like data-dependent shapes and explicit device ownership), we implemented a clean virtual-
routing and monkeypatching solution:
──────

The Monkeypatch: Boolean Indexing ( getitem )

1. The Problem

MLX is designed with an asynchronous graph compilation model where the shape of every array must be statically known at graph-build time.
However, boolean indexing (e.g. x[x < 5] ) is a data-dependent operation: the size of the returned array depends on the number of True values inside
the mask, which isn't known until evaluation. Because of this, MLX arrays natively raise a ValueError: boolean indices are not yet supported when you
try to index them.

2. The Solution

Since the Array API standard guarantees that boolean indexing works, and standard routines (like clip() ) rely on it, we monkeypatched the C++ array
class's indexing method directly:

# In src/array_api_compat/mlx/__init__.py
_old_getitem = mx.array.__getitem__

def _new_getitem(self, item):
    if _has_bool_array(item):
        # Convert array and boolean mask to numpy, index on CPU, and wrap back to MLX
        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

This intercepts indexing calls. If a boolean mask is detected, it falls back to CPU evaluation via NumPy and converts the result back to an MLX array.
If no boolean mask is involved, it passes the call through to MLX's native indexing.

(Note: We did not need to patch setitem because MLX natively supports in-place boolean assignments, e.g. x[mask] = 1.0 .)
──────

The Virtual Tracker: Device and to_device()

1. The Problem

Because MLX arrays reside in unified memory, they do not belong to a device and do not have .device attributes or .to_device() methods.
However, the Array API requires device(to_device(x, dev)) == dev to hold true. If to_device just returned the array (since nothing physically
moves), any subsequent call to check the device would simply return the active system default (usually gpu ), causing device tests to fail when CPU is
selected.

2. The Solution

We implemented a virtual device registry using an id(array) -> Device map in Python:

_mlx_device_map = {}

def _mlx_cleanup(ref_id):
    _mlx_device_map.pop(ref_id, None)

def _set_mlx_device(x, device):
    ref_id = id(x)
    _mlx_device_map[ref_id] = device
    # Prevent memory leaks: clean up when the array is garbage collected
    weakref.finalize(x, _mlx_cleanup, ref_id)

• No eq collisions: We use the integer id(x) rather than a standard WeakKeyDictionary because looking up an array in a WeakKeyDictionary
invokes x == y on keys. Since MLX comparisons return element-wise boolean arrays rather than a single Python boolean, dictionary lookups would raise a
ValueError .
• No memory leaks: Using weakref.finalize automatically pops the ID from the tracking dictionary when the array is garbage collected.

Note: Some help of AI was used for code and pr description writing :)

[Copilot]

Pull request overview

Adds first-class MLX support to array-api-compat, including a wrapped array_api_compat.mlx namespace, MLX-specific linalg/fft shims, and MLX-aware detection + device/to_device behavior in the common helpers, plus accompanying test integration.

Changes:

• Introduces src/array_api_compat/mlx/ (namespace, aliases, inspection info, linalg, fft, typing helpers) and wires it into the build.
• Extends common/_helpers.py with MLX array/namespace detection and a virtual device registry for device()/to_device().
• Adds/updates tests and helpers to include mlx.core in the wrapped-libraries matrix and provides MLX-specific integration tests.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
tests/test_mlx.py	New MLX-focused integration tests for namespace detection, key wrappers, linalg/fft, and inspection info.
tests/test_isdtype.py	Adds MLX-specific spec dtype set for isdtype testing.
tests/test_common.py	Extends common test matrix for MLX detection and adjusts cross-library/asarray-copy expectations.
tests/test_array_namespace.py	Handles mlx.core mapping to array_api_compat.mlx in namespace resolution tests.
tests/_helpers.py	Adds mlx.core to wrapped libraries and updates import helper to skip on top-level package import.
src/array_api_compat/mlx/linalg.py	Adds MLX linalg wrapper with spec-aligned helpers and explicit CPU eigh.
src/array_api_compat/mlx/fft.py	Adds MLX fft wrapper using common FFT helpers for spec behavior.
src/array_api_compat/mlx/_typing.py	Exposes MLX typing aliases (Array/DType/Device).
src/array_api_compat/mlx/_info.py	Implements Array API inspection namespace for MLX capabilities/dtypes/devices.
src/array_api_compat/mlx/_aliases.py	Provides MLX aliases and spec-compat shims (e.g., asarray, iinfo, argsort/sort descending).
src/array_api_compat/mlx/init.py	Creates the MLX compat namespace, loads submodules, and monkeypatches boolean __getitem__.
src/array_api_compat/common/_helpers.py	Adds MLX detection and virtual device mapping for device()/to_device().
pyproject.toml	Adds mlx to the dev dependency group.
meson.build	Includes MLX sources in the build definition.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ev-br]

Thanks for the PR.

We'll need some more context though. I see that it's been discussed for a while, #162, and a suggested avenue was that MLX adds Array API support in its main namespace, similar to JAX.
And there are several discussions in MLX itself, ml-explore/mlx#3514 and references therein.
So what is the upstream status? Upstream adding support is of course the ideal outcome.

Next, what is the current compatibility level of MLX and Array API? I ran the test suite against this PR ($ ARRAY_API_TESTS_MODULE=array_api_compat.mlx pytest array_api_tests/test_creation_functions.py -v), and it fails all around.
Some of these failures might be deficiencies in the test suite itself (cf data-apis/array-api-tests#431 (comment)).

I'd suggest to start from the beginning:

• create an overview document with a list of Array API functions/objects and their MLX analogs (e.g., three columns: an Array API function, a matching MLX function, rough compat status---fundamental issues if any, otherwise name/signature mismatches, arguments, defaults and so on).
• For matching functions, run the corresponding tests from the test suite. Does it pass? If not, what are issues, is the test suite at fault or is MLX incompatible?
• For functions with small differences, would MLX itself accept changes or aliases? (concatenate/concat, acos/arccos etc)
• For situations where the problem is with the test suite, what is missing in it?
• For fundamental problems, what would it take to fix (e.g., monkey-patching the array object we 99/100 won't do here, but let's make it explicit in the design overview)

Overall, I'm pretty sure that the support is doable (at a level comparable with e.g. jax which implements a part of the standard, and pytorch which also implements a part of the standard and has some of the inconsistencies smoothened out in the compat layer). Making it work will take a bit of effort though, spread between the MLX itself, the Array API workgroup if needs be, the test suite and maybe indeed the compat layer. But it all starts with an explicit context and a design overview.

[aaishwarymishra]

My bad, I kind of discussed this with Ralf but you are right, I should get to be more organised. I think most of the methods are implemented now in mlx, accept few and some differences in how the implemented methods behave.

Well the biggest problem right now is no .device() support, and the no float64 support on GPU.

Well float64 can be added in the future but I think we might have to implement the device work around .

Also I will make a nice detailed doc so it would be clearer :)

[rgommers]

A detailed doc sounds quite useful.

Well the biggest problem right now is no .device() support, and the no float64 support on GPU.

There are other libraries with no .device method, which is why SciPy uses the xp_device function as a helper - I don't think this is a major issue.

No float64 on GPU also isn't unique, and will be allowed in the next version of the standard, see the just-merged data-apis/array-api#1005.

[rgommers]

MLX has done a lot of work, my understanding is that all the critical gaps were plugged, and if there are some missing routines, they're open to adding them.