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.