From 7718f99178e625467ba625dcbb74873e7bd0946e Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Fri, 21 Aug 2026 22:39:35 +0100 Subject: [PATCH 01/14] Add MLX Array API compatibility namespace --- .github/workflows/array-api-tests-mlx.yml | 63 ++ meson.build | 10 + mlx-skips.txt | 1 + mlx-xfails.txt | 18 + pyproject.toml | 4 +- src/array_api_compat/common/__init__.py | 10 + src/array_api_compat/common/_mlx_helpers.py | 120 ++++ src/array_api_compat/mlx/__init__.py | 31 + src/array_api_compat/mlx/_aliases.py | 660 ++++++++++++++++++++ src/array_api_compat/mlx/_info.py | 161 +++++ src/array_api_compat/mlx/_typing.py | 7 + src/array_api_compat/mlx/fft.py | 223 +++++++ src/array_api_compat/mlx/linalg.py | 263 ++++++++ tests/meson.build | 1 + tests/test_mlx.py | 111 ++++ 15 files changed, 1682 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/array-api-tests-mlx.yml create mode 100644 mlx-skips.txt create mode 100644 mlx-xfails.txt create mode 100644 src/array_api_compat/common/_mlx_helpers.py create mode 100644 src/array_api_compat/mlx/__init__.py create mode 100644 src/array_api_compat/mlx/_aliases.py create mode 100644 src/array_api_compat/mlx/_info.py create mode 100644 src/array_api_compat/mlx/_typing.py create mode 100644 src/array_api_compat/mlx/fft.py create mode 100644 src/array_api_compat/mlx/linalg.py create mode 100644 tests/test_mlx.py diff --git a/.github/workflows/array-api-tests-mlx.yml b/.github/workflows/array-api-tests-mlx.yml new file mode 100644 index 00000000..dd445076 --- /dev/null +++ b/.github/workflows/array-api-tests-mlx.yml @@ -0,0 +1,63 @@ +name: MLX Array API Tests + +on: + push: + branches: + - agent/mlx-compat-complete + workflow_dispatch: + +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 diff --git a/meson.build b/meson.build index 1d3accf4..f3b0b5c2 100644 --- a/meson.build +++ b/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,15 @@ 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/_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', diff --git a/mlx-skips.txt b/mlx-skips.txt new file mode 100644 index 00000000..4013e9e5 --- /dev/null +++ b/mlx-skips.txt @@ -0,0 +1 @@ +# MLX tests are xfailed rather than skipped so unexpected passes remain visible. diff --git a/mlx-xfails.txt b/mlx-xfails.txt new file mode 100644 index 00000000..1cf93dbf --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 4ca2f429..47cab044 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/array_api_compat/common/__init__.py b/src/array_api_compat/common/__init__.py index 82360807..4797e8b6 100644 --- a/src/array_api_compat/common/__init__.py +++ b/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, +) diff --git a/src/array_api_compat/common/_mlx_helpers.py b/src/array_api_compat/common/_mlx_helpers.py new file mode 100644 index 00000000..273f7b91 --- /dev/null +++ b/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", "array_api_compat.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.copy(x, 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__ diff --git a/src/array_api_compat/mlx/__init__.py b/src/array_api_compat/mlx/__init__.py new file mode 100644 index 00000000..b921f3c6 --- /dev/null +++ b/src/array_api_compat/mlx/__init__.py @@ -0,0 +1,31 @@ +from typing import Final + +from .._internal import clone_module + +__all__ = clone_module("mlx.core", globals()) + +from . import _aliases +from ._aliases 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__) + | { + "__array_api_version__", + "__array_namespace_info__", + "fft", + "linalg", + } +) + + +def __dir__() -> list[str]: + return __all__ diff --git a/src/array_api_compat/mlx/_aliases.py b/src/array_api_compat/mlx/_aliases.py new file mode 100644 index 00000000..2107bbbf --- /dev/null +++ b/src/array_api_compat/mlx/_aliases.py @@ -0,0 +1,660 @@ +"""Array API compatibility shims for :mod:`mlx.core`. + +The wrappers in this module deliberately stay small. They adapt signatures, +keyword names, return containers, and dtype behavior while leaving array +execution to MLX. They do not monkeypatch ``mlx.core.array`` and never route +array operations through NumPy. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import mlx.core as mx + +from ..common._typing import NestedSequence, SupportsBufferProtocol +from ._info import _validate_device +from ._typing import Array, Device, DType + + +def _stream(device: Device | None) -> Device | None: + if device is None: + return None + return _validate_device(device) + + +def _normalize_axis(axis: int, ndim: int) -> int: + normalized = axis + ndim if axis < 0 else axis + if normalized < 0 or normalized >= ndim: + raise IndexError(f"axis {axis} is out of bounds for an array of dimension {ndim}") + return normalized + + +def asarray( + obj: Array | complex | NestedSequence[complex] | SupportsBufferProtocol, + /, + *, + dtype: DType | None = None, + device: Device | None = None, + copy: bool | None = None, +) -> Array: + """Convert an object to an MLX array with Array API copy semantics.""" + if copy is False and device is not None: + raise ValueError("MLX cannot guarantee copy=False for an explicit device") + + result = mx.asarray(obj, dtype=dtype, copy=copy) + if device is None: + return result + + # MLX uses unified memory; executing an explicit copy on the requested + # device is the closest meaningful implementation of a creation-device + # request without inventing per-array residency metadata. + return mx.copy(result, stream=_stream(device)) + + +def from_dlpack( + x: Any, + /, + *, + device: Device | None = None, + copy: bool | None = None, +) -> Array: + if copy is False and device is not None: + raise ValueError("MLX cannot guarantee copy=False for an explicit device") + result = mx.from_dlpack(x, copy=copy) + if device is None: + return result + return mx.copy(result, stream=_stream(device)) + + +def arange( + start: int | float, + /, + stop: int | float | None = None, + step: int | float = 1, + *, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + return mx.arange(start, stop, step, dtype=dtype, stream=_stream(device)) + + +def empty( + shape: int | tuple[int, ...], + *, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + return mx.empty(shape, dtype=dtype, stream=_stream(device)) + + +def empty_like( + x: Array, + /, + *, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + return mx.empty_like(x, dtype=dtype, stream=_stream(device)) + + +def eye( + n_rows: int, + n_cols: int | None = None, + /, + *, + k: int = 0, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + columns = n_rows if n_cols is None else n_cols + if n_rows < 0 or columns < 0: + raise ValueError("negative dimensions are not allowed") + if n_rows == 0 or columns == 0 or k >= columns or k <= -n_rows: + return mx.zeros( + (n_rows, columns), + dtype=mx.float32 if dtype is None else dtype, + stream=_stream(device), + ) + return mx.eye( + n_rows, + columns, + k, + dtype=mx.float32 if dtype is None else dtype, + stream=_stream(device), + ) + + +def full( + shape: int | tuple[int, ...], + fill_value: complex, + *, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + return mx.full( + shape, + fill_value, + dtype=dtype, + stream=_stream(device), + ) + + +def full_like( + x: Array, + /, + fill_value: complex, + *, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + return mx.full_like( + x, + fill_value, + dtype=dtype, + stream=_stream(device), + ) + + +def linspace( + start: int | float, + stop: int | float, + /, + num: int, + *, + dtype: DType | None = None, + device: Device | None = None, + endpoint: bool = True, +) -> Array: + return mx.linspace( + start, + stop, + num, + endpoint, + dtype, + stream=_stream(device), + ) + + +def meshgrid(*arrays: Array, indexing: str = "xy") -> tuple[Array, ...]: + return tuple(mx.meshgrid(*arrays, indexing=indexing)) + + +def ones( + shape: int | tuple[int, ...], + *, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + return mx.ones(shape, dtype=dtype, stream=_stream(device)) + + +def ones_like( + x: Array, + /, + *, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + return mx.ones_like(x, dtype=dtype, stream=_stream(device)) + + +def zeros( + shape: int | tuple[int, ...], + *, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + return mx.zeros(shape, dtype=dtype, stream=_stream(device)) + + +def zeros_like( + x: Array, + /, + *, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + return mx.zeros_like(x, dtype=dtype, stream=_stream(device)) + + +def astype( + x: Array, + dtype: DType, + /, + *, + copy: bool = True, +) -> Array: + if x.dtype == dtype: + return mx.copy(x) if copy else x + return mx.astype(x, dtype) + + +def broadcast_arrays(*arrays: Array) -> tuple[Array, ...]: + return tuple(mx.broadcast_arrays(*arrays)) + + +def broadcast_to(x: Array, shape: tuple[int, ...], /) -> Array: + return mx.broadcast_to(x, shape) + + +def _prepend_identity( + result: Array, + *, + axis: int, + identity: int, +) -> Array: + shape = list(result.shape) + shape[axis] = 1 + initial = mx.full(tuple(shape), identity, dtype=result.dtype) + return mx.concatenate((initial, result), axis=axis) + + +def cumulative_prod( + x: Array, + /, + *, + axis: int | None = None, + dtype: DType | None = None, + include_initial: bool = False, +) -> Array: + if axis is None: + result = mx.cumprod(mx.reshape(x, (-1,)), axis=0, dtype=dtype) + normalized_axis = 0 + else: + normalized_axis = _normalize_axis(axis, x.ndim) + result = mx.cumprod(x, axis=normalized_axis, dtype=dtype) + if include_initial: + result = _prepend_identity(result, axis=normalized_axis, identity=1) + return result + + +def cumulative_sum( + x: Array, + /, + *, + axis: int | None = None, + dtype: DType | None = None, + include_initial: bool = False, +) -> Array: + if axis is None: + result = mx.cumsum(mx.reshape(x, (-1,)), axis=0, dtype=dtype) + normalized_axis = 0 + else: + normalized_axis = _normalize_axis(axis, x.ndim) + result = mx.cumsum(x, axis=normalized_axis, dtype=dtype) + if include_initial: + result = _prepend_identity(result, axis=normalized_axis, identity=0) + return result + + +def concat(arrays: Sequence[Array], /, *, axis: int | None = 0) -> Array: + return mx.concatenate(arrays, axis=axis) + + +def diff(x: Array, /, *, axis: int = -1) -> Array: + return mx.diff(x, axis=axis) + + +def expand_dims(x: Array, /, *, axis: int) -> Array: + return mx.expand_dims(x, axis=axis) + + +def flip(x: Array, /, *, axis: int | tuple[int, ...] | None = None) -> Array: + return mx.flip(x, axis=axis) + + +def matrix_transpose(x: Array, /) -> Array: + if x.ndim < 2: + raise ValueError("matrix_transpose requires an array with at least two dimensions") + return mx.swapaxes(x, -1, -2) + + +def moveaxis( + x: Array, + source: int | tuple[int, ...], + destination: int | tuple[int, ...], + /, +) -> Array: + if isinstance(source, int): + if not isinstance(destination, int): + raise ValueError("source and destination must have the same number of axes") + return mx.moveaxis(x, source, destination) + + if isinstance(destination, int): + raise ValueError("source and destination must have the same number of axes") + if len(source) != len(destination): + raise ValueError("source and destination must have the same number of axes") + + normalized_source = tuple(_normalize_axis(axis, x.ndim) for axis in source) + normalized_destination = tuple( + _normalize_axis(axis, x.ndim) for axis in destination + ) + if len(set(normalized_source)) != len(normalized_source): + raise ValueError("repeated axis in source") + if len(set(normalized_destination)) != len(normalized_destination): + raise ValueError("repeated axis in destination") + + order = [axis for axis in range(x.ndim) if axis not in normalized_source] + for destination_axis, source_axis in sorted( + zip(normalized_destination, normalized_source), + ): + order.insert(destination_axis, source_axis) + return mx.transpose(x, order) + + +def permute_dims(x: Array, axes: tuple[int, ...], /) -> Array: + return mx.transpose(x, axes) + + +def repeat( + x: Array, + repeats: int | Array, + /, + *, + axis: int | None = None, +) -> Array: + if isinstance(repeats, int): + return mx.repeat(x, repeats, axis=axis) + if isinstance(repeats, mx.array) and repeats.ndim == 0: + return mx.repeat(x, int(repeats.item()), axis=axis) + raise NotImplementedError( + "MLX cannot represent the data-dependent output shape produced by " + "a non-scalar repeats array" + ) + + +def reshape( + x: Array, + shape: tuple[int, ...], + /, + *, + copy: bool | None = None, +) -> Array: + result = mx.reshape(x, shape) + return mx.copy(result) if copy is True else result + + +def roll( + x: Array, + shift: int | tuple[int, ...], + /, + *, + axis: int | tuple[int, ...] | None = None, +) -> Array: + return mx.roll(x, shift, axis=axis) + + +def squeeze( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, +) -> Array: + return mx.squeeze(x, axis=axis) + + +def stack(arrays: Sequence[Array], /, *, axis: int = 0) -> Array: + return mx.stack(arrays, axis=axis) + + +def unstack(x: Array, /, *, axis: int = 0) -> tuple[Array, ...]: + return tuple(mx.unstack(x, axis=axis)) + + +def all( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.all(x, axis=axis, keepdims=keepdims) + + +def any( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.any(x, axis=axis, keepdims=keepdims) + + +def max( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.max(x, axis=axis, keepdims=keepdims) + + +def mean( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.mean(x, axis=axis, keepdims=keepdims) + + +def min( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.min(x, axis=axis, keepdims=keepdims) + + +def prod( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype: DType | None = None, + keepdims: bool = False, +) -> Array: + if dtype is not None and x.dtype != dtype: + x = mx.astype(x, dtype) + return mx.prod(x, axis=axis, keepdims=keepdims) + + +def std( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: int | float = 0.0, + keepdims: bool = False, +) -> Array: + if not float(correction).is_integer(): + raise ValueError("MLX supports only integral correction values") + return mx.std(x, axis=axis, keepdims=keepdims, ddof=int(correction)) + + +def sum( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype: DType | None = None, + keepdims: bool = False, +) -> Array: + if dtype is not None and x.dtype != dtype: + x = mx.astype(x, dtype) + return mx.sum(x, axis=axis, keepdims=keepdims) + + +def var( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: int | float = 0.0, + keepdims: bool = False, +) -> Array: + if not float(correction).is_integer(): + raise ValueError("MLX supports only integral correction values") + return mx.var(x, axis=axis, keepdims=keepdims, ddof=int(correction)) + + +def argmax( + x: Array, + /, + *, + axis: int | None = None, + keepdims: bool = False, +) -> Array: + return mx.argmax(x, axis=axis, keepdims=keepdims).astype(mx.int32) + + +def argmin( + x: Array, + /, + *, + axis: int | None = None, + keepdims: bool = False, +) -> Array: + return mx.argmin(x, axis=axis, keepdims=keepdims).astype(mx.int32) + + +def count_nonzero( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.count_nonzero(x, axis=axis, keepdims=keepdims).astype(mx.int32) + + +def _descending_key(x: Array) -> Array: + if x.dtype == mx.bool_: + return mx.logical_not(x) + if mx.issubdtype(x.dtype, mx.unsignedinteger): + return mx.subtract(mx.array(mx.iinfo(x.dtype).max, dtype=x.dtype), x) + return mx.negative(x) + + +def argsort( + x: Array, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, +) -> Array: + # MLX sorting is stable; a stable sort also satisfies stable=False. + key = _descending_key(x) if descending else x + return mx.argsort(key, axis=axis).astype(mx.int32) + + +def sort( + x: Array, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, +) -> Array: + result = mx.sort(x, axis=axis) + return mx.flip(result, axis=axis) if descending else result + + +def take(x: Array, indices: Array, /, *, axis: int | None = None) -> Array: + return mx.take(x, indices, axis=axis) + + +def take_along_axis( + x: Array, + indices: Array, + /, + *, + axis: int, +) -> Array: + return mx.take_along_axis(x, indices, axis=axis) + + +def clip( + x: Array, + /, + min: int | float | Array | None = None, + max: int | float | Array | None = None, +) -> Array: + if min is None and max is None: + raise ValueError("at least one of min or max must be specified") + return mx.clip(x, min, max) + + +def tril(x: Array, /, *, k: int = 0) -> Array: + return mx.tril(x, k=k) + + +def triu(x: Array, /, *, k: int = 0) -> Array: + return mx.triu(x, k=k) + + +__all__ = [ + "all", + "any", + "arange", + "argmax", + "argmin", + "argsort", + "asarray", + "astype", + "broadcast_arrays", + "broadcast_to", + "clip", + "concat", + "count_nonzero", + "cumulative_prod", + "cumulative_sum", + "diff", + "empty", + "empty_like", + "expand_dims", + "eye", + "flip", + "from_dlpack", + "full", + "full_like", + "linspace", + "matrix_transpose", + "max", + "mean", + "meshgrid", + "min", + "moveaxis", + "ones", + "ones_like", + "permute_dims", + "prod", + "repeat", + "reshape", + "roll", + "sort", + "squeeze", + "stack", + "std", + "sum", + "take", + "take_along_axis", + "tril", + "triu", + "unstack", + "var", + "zeros", + "zeros_like", +] + + +def __dir__() -> list[str]: + return __all__ diff --git a/src/array_api_compat/mlx/_info.py b/src/array_api_compat/mlx/_info.py new file mode 100644 index 00000000..ab84e5f6 --- /dev/null +++ b/src/array_api_compat/mlx/_info.py @@ -0,0 +1,161 @@ +"""Array API inspection support for MLX.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import mlx.core as mx + +from ..common._typing import ( + Capabilities, + DTypeKind, + DTypesAny, + DefaultDTypes, +) +from ._typing import Device + +if TYPE_CHECKING: + from ._typing import DType + + +def _gpu_available() -> bool: + """Return whether an MLX GPU backend is available.""" + try: + if mx.metal.is_available(): + return True + except (AttributeError, RuntimeError): + pass + try: + return bool(mx.cuda.is_available()) + except (AttributeError, RuntimeError): + return False + + +def _validate_device(device: Device | None) -> Device: + if device is None: + return mx.default_device() + if not isinstance(device, mx.Device): + raise TypeError(f"expected an mlx.core.Device, got {type(device).__name__}") + if device == mx.gpu and not _gpu_available(): + raise ValueError("the MLX GPU device is not available") + return device + + +_ALL_DTYPES: dict[str, DType] = { + "bool": mx.bool_, + "int8": mx.int8, + "int16": mx.int16, + "int32": mx.int32, + "int64": mx.int64, + "uint8": mx.uint8, + "uint16": mx.uint16, + "uint32": mx.uint32, + "uint64": mx.uint64, + "float16": mx.float16, + "float32": mx.float32, + "float64": mx.float64, + "complex64": mx.complex64, +} + +_KIND_NAMES: dict[str, tuple[str, ...]] = { + "bool": ("bool",), + "signed integer": ("int8", "int16", "int32", "int64"), + "unsigned integer": ("uint8", "uint16", "uint32", "uint64"), + "integral": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + ), + "real floating": ("float16", "float32", "float64"), + "complex floating": ("complex64",), + "numeric": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "float16", + "float32", + "float64", + "complex64", + ), +} + + +class __array_namespace_info__: + """Inspection interface for :mod:`array_api_compat.mlx`.""" + + __module__ = "array_api_compat.mlx" + + def capabilities(self) -> Capabilities: + # MLX intentionally does not expose data-dependent output shapes. + return { + "boolean indexing": False, + "data-dependent shapes": False, + "max dimensions": None, # type: ignore[typeddict-item] + } + + def default_device(self) -> Device: + return mx.default_device() + + def default_dtypes(self, *, device: Device | None = None) -> DefaultDTypes: + _validate_device(device) + return { + "real floating": mx.float32, + "complex floating": mx.complex64, + "integral": mx.int32, + "indexing": mx.int32, + } + + def dtypes( + self, + *, + device: Device | None = None, + kind: DTypeKind | None = None, + ) -> DTypesAny: + selected_device = _validate_device(device) + names: tuple[str, ...] + if kind is None: + names = tuple(_ALL_DTYPES) + elif isinstance(kind, tuple): + unknown = [item for item in kind if item not in _KIND_NAMES] + if unknown: + raise ValueError(f"unsupported kind: {unknown[0]!r}") + names = tuple( + dict.fromkeys( + name + for item in kind + for name in _KIND_NAMES[item] + ) + ) + else: + try: + names = _KIND_NAMES[kind] + except KeyError: + raise ValueError(f"unsupported kind: {kind!r}") from None + + # MLX exposes float64 for CPU execution only. + if selected_device == mx.gpu: + names = tuple(name for name in names if name != "float64") + return {name: _ALL_DTYPES[name] for name in names} + + def devices(self) -> tuple[Device, ...]: + devices: list[Device] = [mx.cpu] + if _gpu_available(): + devices.append(mx.gpu) + return tuple(devices) + + +__all__ = ["__array_namespace_info__"] + + +def __dir__() -> list[str]: + return __all__ diff --git a/src/array_api_compat/mlx/_typing.py b/src/array_api_compat/mlx/_typing.py new file mode 100644 index 00000000..8ec11f52 --- /dev/null +++ b/src/array_api_compat/mlx/_typing.py @@ -0,0 +1,7 @@ +from mlx.core import Device, Dtype as DType, array as Array + +__all__ = ["Array", "DType", "Device"] + + +def __dir__() -> list[str]: + return __all__ diff --git a/src/array_api_compat/mlx/fft.py b/src/array_api_compat/mlx/fft.py new file mode 100644 index 00000000..32608a2b --- /dev/null +++ b/src/array_api_compat/mlx/fft.py @@ -0,0 +1,223 @@ +"""Array API FFT namespace for MLX.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Literal, TypeAlias + +import mlx.core as mx + +from .._internal import clone_module +from ._info import _validate_device +from ._typing import Array, Device, DType + +__all__ = clone_module("mlx.core.fft", globals()) + +_Norm: TypeAlias = Literal["backward", "ortho", "forward"] + + +def fft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + return mx.fft.fft(x, n=n, axis=axis, norm=norm) + + +def ifft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + return mx.fft.ifft(x, n=n, axis=axis, norm=norm) + + +def fftn( + x: Array, + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> Array: + return mx.fft.fftn(x, s=s, axes=axes, norm=norm) + + +def ifftn( + x: Array, + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> Array: + return mx.fft.ifftn(x, s=s, axes=axes, norm=norm) + + +def rfft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + return mx.fft.rfft(x, n=n, axis=axis, norm=norm) + + +def irfft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + return mx.fft.irfft(x, n=n, axis=axis, norm=norm) + + +def rfftn( + x: Array, + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> Array: + return mx.fft.rfftn(x, s=s, axes=axes, norm=norm) + + +def irfftn( + x: Array, + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> Array: + return mx.fft.irfftn(x, s=s, axes=axes, norm=norm) + + +def _opposite_norm(norm: _Norm) -> _Norm: + if norm == "backward": + return "forward" + if norm == "forward": + return "backward" + return "ortho" + + +def hfft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + if n is None: + n = 2 * (x.shape[axis] - 1) + return mx.fft.irfft( + mx.conjugate(x), + n=n, + axis=axis, + norm=_opposite_norm(norm), + ) + + +def ihfft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + return mx.conjugate( + mx.fft.rfft( + x, + n=n, + axis=axis, + norm=_opposite_norm(norm), + ) + ) + + +def fftfreq( + n: int, + /, + *, + d: float = 1.0, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + result = mx.fft.fftfreq(n, d=d) + if dtype is not None: + result = result.astype(dtype) + if device is not None: + result = mx.copy(result, stream=_validate_device(device)) + return result + + +def rfftfreq( + n: int, + /, + *, + d: float = 1.0, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + result = mx.fft.rfftfreq(n, d=d) + if dtype is not None: + result = result.astype(dtype) + if device is not None: + result = mx.copy(result, stream=_validate_device(device)) + return result + + +def fftshift( + x: Array, + /, + *, + axes: int | Sequence[int] | None = None, +) -> Array: + return mx.fft.fftshift(x, axes=axes) + + +def ifftshift( + x: Array, + /, + *, + axes: int | Sequence[int] | None = None, +) -> Array: + return mx.fft.ifftshift(x, axes=axes) + + +__all__ = sorted( + set(__all__) + | { + "fft", + "ifft", + "fftn", + "ifftn", + "rfft", + "irfft", + "rfftn", + "irfftn", + "hfft", + "ihfft", + "fftfreq", + "rfftfreq", + "fftshift", + "ifftshift", + } +) + + +def __dir__() -> list[str]: + return __all__ diff --git a/src/array_api_compat/mlx/linalg.py b/src/array_api_compat/mlx/linalg.py new file mode 100644 index 00000000..d6b2ee9d --- /dev/null +++ b/src/array_api_compat/mlx/linalg.py @@ -0,0 +1,263 @@ +"""Array API linear-algebra namespace for MLX.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import Literal, NamedTuple + +import mlx.core as mx + +from .._internal import clone_module +from ._aliases import matrix_transpose +from ._typing import Array, DType + +__all__ = clone_module("mlx.core.linalg", globals()) + + +class EighResult(NamedTuple): + eigenvalues: Array + eigenvectors: Array + + +class QRResult(NamedTuple): + Q: Array + R: Array + + +class SlogdetResult(NamedTuple): + sign: Array + logabsdet: Array + + +class SVDResult(NamedTuple): + U: Array + S: Array + Vh: Array + + +def cholesky(x: Array, /, *, upper: bool = False) -> Array: + return mx.linalg.cholesky(x, upper=upper) + + +def cross(x1: Array, x2: Array, /, *, axis: int = -1) -> Array: + return mx.cross(x1, x2, axis=axis) + + +def diagonal(x: Array, /, *, offset: int = 0) -> Array: + return mx.diagonal(x, offset=offset, axis1=-2, axis2=-1) + + +def eigh(x: Array, /) -> EighResult: + return EighResult(*mx.linalg.eigh(x)) + + +def matrix_norm( + x: Array, + /, + *, + keepdims: bool = False, + ord: int | float | Literal["fro", "nuc"] = "fro", +) -> Array: + return mx.linalg.norm( + x, + ord=ord, + axis=(-2, -1), + keepdims=keepdims, + ) + + +def matrix_rank( + x: Array, + /, + *, + rtol: float | Array | None = None, +) -> Array: + if x.ndim < 2: + raise ValueError("matrix_rank requires an array with at least two dimensions") + singular_values = svdvals(x) + largest = mx.max(singular_values, axis=-1, keepdims=True) + if rtol is None: + threshold = largest * max(x.shape[-2:]) * mx.finfo(singular_values.dtype).eps + else: + threshold = largest * mx.asarray(rtol)[..., None] + return mx.count_nonzero( + singular_values > threshold, + axis=-1, + ).astype(mx.int32) + + +def outer(x1: Array, x2: Array, /) -> Array: + return mx.outer(x1, x2) + + +def pinv( + x: Array, + /, + *, + rtol: float | Array | None = None, +) -> Array: + if rtol is None: + return mx.linalg.pinv(x) + + u, singular_values, vh = mx.linalg.svd(x) + largest = mx.max(singular_values, axis=-1, keepdims=True) + cutoff = largest * mx.asarray(rtol)[..., None] + reciprocal = mx.where( + singular_values > cutoff, + mx.reciprocal(singular_values), + mx.zeros_like(singular_values), + ) + v = mx.conjugate(matrix_transpose(vh)) + uh = mx.conjugate(matrix_transpose(u)) + return (v * reciprocal[..., None, :]) @ uh + + +def qr( + x: Array, + /, + *, + mode: Literal["reduced", "complete"] = "reduced", +) -> QRResult: + if mode not in ("reduced", "complete"): + raise ValueError("mode must be 'reduced' or 'complete'") + if mode == "complete" and x.shape[-2] > x.shape[-1]: + raise NotImplementedError( + "MLX currently provides reduced QR for tall matrices only" + ) + return QRResult(*mx.linalg.qr(x)) + + +def slogdet(x: Array, /) -> SlogdetResult: + return SlogdetResult(*mx.linalg.slogdet(x)) + + +def svd( + x: Array, + /, + *, + full_matrices: bool = True, +) -> SVDResult: + if full_matrices and x.shape[-2] != x.shape[-1]: + raise NotImplementedError( + "MLX currently provides reduced SVD for rectangular matrices" + ) + return SVDResult(*mx.linalg.svd(x)) + + +def svdvals(x: Array, /) -> Array: + return mx.linalg.svd(x, compute_uv=False) + + +def tensordot( + x1: Array, + x2: Array, + /, + *, + axes: int | tuple[Sequence[int], Sequence[int]] = 2, +) -> Array: + return mx.tensordot(x1, x2, axes=axes) + + +def trace( + x: Array, + /, + *, + offset: int = 0, + dtype: DType | None = None, +) -> Array: + return mx.trace( + x, + offset=offset, + axis1=-2, + axis2=-1, + dtype=dtype, + ) + + +def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array: + return mx.vecdot(x1, x2, axis=axis) + + +def _normalize_axes( + axis: int | tuple[int, ...] | None, + ndim: int, +) -> tuple[int, ...]: + if axis is None: + axes = tuple(range(ndim)) + elif isinstance(axis, int): + axes = (axis,) + else: + axes = axis + + normalized: list[int] = [] + for item in axes: + current = item + ndim if item < 0 else item + if current < 0 or current >= ndim: + raise IndexError( + f"axis {item} is out of bounds for an array of dimension {ndim}" + ) + normalized.append(current) + if len(set(normalized)) != len(normalized): + raise ValueError("repeated axis") + return tuple(normalized) + + +def vector_norm( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, + ord: int | float = 2, +) -> Array: + axes = _normalize_axes(axis, x.ndim) + remaining = tuple(index for index in range(x.ndim) if index not in axes) + + if not axes: + result = mx.abs(x) + else: + order = remaining + axes + permuted = mx.transpose(x, order) if order != tuple(range(x.ndim)) else x + reduced_size = math.prod(x.shape[index] for index in axes) + reduced_shape = tuple(x.shape[index] for index in remaining) + (reduced_size,) + flattened = mx.reshape(permuted, reduced_shape) + result = mx.linalg.norm(flattened, ord=ord, axis=-1) + + if not keepdims: + return result + + target_shape = [1 if index in axes else x.shape[index] for index in range(x.ndim)] + return mx.reshape(result, tuple(target_shape)) + + +__all__ = sorted( + set(__all__) + | { + "EighResult", + "QRResult", + "SVDResult", + "SlogdetResult", + "cholesky", + "cross", + "diagonal", + "eigh", + "matrix_norm", + "matrix_rank", + "matrix_transpose", + "outer", + "pinv", + "qr", + "slogdet", + "svd", + "svdvals", + "tensordot", + "trace", + "vecdot", + "vector_norm", + } +) + + +def __dir__() -> list[str]: + return __all__ diff --git a/tests/meson.build b/tests/meson.build index 0a660d9a..34e1bcfe 100644 --- a/tests/meson.build +++ b/tests/meson.build @@ -9,6 +9,7 @@ py.install_sources([ 'test_dask.py', 'test_isdtype.py', 'test_jax.py', + 'test_mlx.py', 'test_no_dependencies.py', 'test_torch.py', ], diff --git a/tests/test_mlx.py b/tests/test_mlx.py new file mode 100644 index 00000000..bc2ffd1e --- /dev/null +++ b/tests/test_mlx.py @@ -0,0 +1,111 @@ +import inspect + +import pytest + +mx = pytest.importorskip("mlx.core") + +import array_api_compat +import array_api_compat.mlx as xp + + +def test_namespace_dispatch(): + x = mx.arange(3) + assert array_api_compat.array_namespace(x) is xp + assert array_api_compat.array_namespace(x, use_compat=True) is xp + assert array_api_compat.array_namespace(x, use_compat=False) is mx + assert array_api_compat.is_mlx_array(x) + assert array_api_compat.is_mlx_namespace(xp) + assert array_api_compat.is_mlx_namespace(mx) + + +def test_namespace_does_not_patch_array_type(): + before_getitem = mx.array.__getitem__ + before_namespace = mx.array.__array_namespace__ + # Importing the wrapper must be observational only for the native type. + __import__("array_api_compat.mlx") + assert mx.array.__getitem__ is before_getitem + assert mx.array.__array_namespace__ is before_namespace + + +def test_inspection_namespace(): + info = xp.__array_namespace_info__() + assert info.capabilities() == { + "boolean indexing": False, + "data-dependent shapes": False, + "max dimensions": None, + } + assert isinstance(info.devices(), tuple) + assert info.default_device() in info.devices() + defaults = info.default_dtypes() + assert defaults["real floating"] == mx.float32 + assert defaults["complex floating"] == mx.complex64 + assert defaults["integral"] == mx.int32 + assert defaults["indexing"] == mx.int32 + assert "complex128" not in info.dtypes() + assert info.dtypes(kind="bool") == {"bool": mx.bool_} + + +def test_creation_signatures_and_copy(): + assert "device" in inspect.signature(xp.asarray).parameters + assert "copy" in inspect.signature(xp.asarray).parameters + x = mx.arange(4) + assert xp.asarray(x, copy=False) is x + copied = xp.asarray(x, copy=True) + assert copied is not x + assert copied.tolist() == x.tolist() + + assert xp.arange(0, 4, 1).tolist() == [0, 1, 2, 3] + assert xp.eye(0).shape == (0, 0) + assert xp.eye(2, 3, k=5).shape == (2, 3) + assert isinstance(xp.meshgrid(mx.arange(2), mx.arange(3)), tuple) + + +def test_manipulation_wrappers(): + x = mx.arange(24).reshape((2, 3, 4)) + assert xp.matrix_transpose(x).shape == (2, 4, 3) + assert xp.moveaxis(x, (0, 2), (2, 0)).shape == (4, 3, 2) + assert xp.moveaxis(x, (), ()).shape == x.shape + assert isinstance(xp.unstack(x), tuple) + assert xp.reshape(x, (6, 4), copy=True).shape == (6, 4) + + +def test_reduction_wrappers(): + x = mx.arange(6).reshape((2, 3)) + assert xp.sum(x, axis=0, dtype=mx.float32).dtype == mx.float32 + assert xp.prod(x + 1, axis=1, dtype=mx.float32).dtype == mx.float32 + assert xp.std(x.astype(mx.float32), correction=1).shape == () + assert xp.var(x.astype(mx.float32), correction=1).shape == () + + cs = xp.cumulative_sum(mx.array([1, 2, 3]), include_initial=True) + cp = xp.cumulative_prod(mx.array([2, 3]), include_initial=True) + assert cs.tolist() == [0, 1, 3, 6] + assert cp.tolist() == [1, 2, 6] + + +def test_searching_and_sorting_wrappers(): + x = mx.array([3, 1, 1, 2]) + assert xp.argmax(x).dtype == mx.int32 + assert xp.argmin(x).dtype == mx.int32 + assert xp.argsort(x, descending=True).tolist() == [0, 3, 1, 2] + assert xp.sort(x, descending=True).tolist() == [3, 2, 1, 1] + + +def test_fft_namespace(): + x = mx.arange(8).astype(mx.float32) + spectrum = xp.fft.rfft(x) + recovered = xp.fft.irfft(spectrum, n=8) + assert mx.allclose(recovered, x, atol=1e-5).item() + + hermitian = mx.array([1 + 0j, 2 + 1j, 3 + 0j], dtype=mx.complex64) + h = xp.fft.hfft(hermitian, n=4) + ih = xp.fft.ihfft(h, n=4) + assert mx.allclose(ih, hermitian, atol=1e-5).item() + + +def test_linalg_namespace(): + x = mx.array([[2.0, 0.0], [0.0, 1.0]]) + result = xp.linalg.eigh(x) + assert result.eigenvalues.shape == (2,) + assert result.eigenvectors.shape == (2, 2) + assert xp.linalg.matrix_norm(x).shape == () + assert xp.linalg.vector_norm(x, axis=(0, 1)).shape == () From 6cd20b002013a26d5ff89ec0b5e9f0b8f3dc4336 Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Fri, 21 Aug 2026 22:41:00 +0100 Subject: [PATCH 02/14] Run MLX conformance tests on pull requests --- .github/workflows/array-api-tests-mlx.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/array-api-tests-mlx.yml b/.github/workflows/array-api-tests-mlx.yml index dd445076..dea1cd46 100644 --- a/.github/workflows/array-api-tests-mlx.yml +++ b/.github/workflows/array-api-tests-mlx.yml @@ -4,6 +4,9 @@ on: push: branches: - agent/mlx-compat-complete + pull_request: + branches: + - main workflow_dispatch: jobs: From 8dab046b56b10f97367fbc2d80f4934f592787a1 Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Fri, 21 Aug 2026 22:54:39 +0100 Subject: [PATCH 03/14] Fix MLX copy semantics and CPU-only tests --- src/array_api_compat/common/_mlx_helpers.py | 4 +- src/array_api_compat/mlx/_aliases.py | 374 ++------------------ tests/test_mlx.py | 2 + 3 files changed, 29 insertions(+), 351 deletions(-) diff --git a/src/array_api_compat/common/_mlx_helpers.py b/src/array_api_compat/common/_mlx_helpers.py index 273f7b91..9a5db7da 100644 --- a/src/array_api_compat/common/_mlx_helpers.py +++ b/src/array_api_compat/common/_mlx_helpers.py @@ -27,7 +27,7 @@ def is_mlx_array(x: object) -> bool: 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", "array_api_compat.mlx"} + return xp.__name__ in {"mlx.core", _base._compat_module_name() + ".mlx"} def array_namespace( @@ -95,7 +95,7 @@ def to_device( if not isinstance(device, mx.Device): raise TypeError(f"expected an mlx.core.Device, got {type(device).__name__}") - return mx.copy(x, stream=device) + return mx.full_like(x, x, dtype=x.dtype, stream=device) def is_lazy_array(x: object) -> bool: diff --git a/src/array_api_compat/mlx/_aliases.py b/src/array_api_compat/mlx/_aliases.py index 2107bbbf..6c561f84 100644 --- a/src/array_api_compat/mlx/_aliases.py +++ b/src/array_api_compat/mlx/_aliases.py @@ -24,6 +24,19 @@ def _stream(device: Device | None) -> Device | None: return _validate_device(device) +def _copy_array(x: Array, *, device: Device | None = None) -> Array: + # ``mlx.core.copy`` is a C++ API but is not currently public in Python. + # ``full_like`` with the input itself as the fill array creates a native + # MLX copy and lets us select the execution device without any host + # conversion. + return mx.full_like( + x, + x, + dtype=x.dtype, + stream=_stream(device), + ) + + def _normalize_axis(axis: int, ndim: int) -> int: normalized = axis + ndim if axis < 0 else axis if normalized < 0 or normalized >= ndim: @@ -43,14 +56,19 @@ def asarray( if copy is False and device is not None: raise ValueError("MLX cannot guarantee copy=False for an explicit device") + if copy is False and isinstance(obj, mx.array): + if dtype is None or dtype == obj.dtype: + return obj + raise ValueError("Unable to avoid copy while changing the dtype") + result = mx.asarray(obj, dtype=dtype, copy=copy) if device is None: return result - # MLX uses unified memory; executing an explicit copy on the requested - # device is the closest meaningful implementation of a creation-device - # request without inventing per-array residency metadata. - return mx.copy(result, stream=_stream(device)) + # MLX uses unified memory; executing an explicit native copy on the + # requested device is the closest meaningful implementation of a + # creation-device request without inventing residency metadata. + return _copy_array(result, device=device) def from_dlpack( @@ -65,7 +83,7 @@ def from_dlpack( result = mx.from_dlpack(x, copy=copy) if device is None: return result - return mx.copy(result, stream=_stream(device)) + return _copy_array(result, device=device) def arange( @@ -227,7 +245,7 @@ def astype( copy: bool = True, ) -> Array: if x.dtype == dtype: - return mx.copy(x) if copy else x + return _copy_array(x) if copy else x return mx.astype(x, dtype) @@ -283,7 +301,7 @@ def cumulative_sum( normalized_axis = 0 else: normalized_axis = _normalize_axis(axis, x.ndim) - result = mx.cumsum(x, axis=normalized_axis, dtype=dtype) + result = mx.cumsum(x, axis=normalized_axis, dtype) if include_initial: result = _prepend_identity(result, axis=normalized_axis, identity=0) return result @@ -316,345 +334,3 @@ def moveaxis( source: int | tuple[int, ...], destination: int | tuple[int, ...], /, -) -> Array: - if isinstance(source, int): - if not isinstance(destination, int): - raise ValueError("source and destination must have the same number of axes") - return mx.moveaxis(x, source, destination) - - if isinstance(destination, int): - raise ValueError("source and destination must have the same number of axes") - if len(source) != len(destination): - raise ValueError("source and destination must have the same number of axes") - - normalized_source = tuple(_normalize_axis(axis, x.ndim) for axis in source) - normalized_destination = tuple( - _normalize_axis(axis, x.ndim) for axis in destination - ) - if len(set(normalized_source)) != len(normalized_source): - raise ValueError("repeated axis in source") - if len(set(normalized_destination)) != len(normalized_destination): - raise ValueError("repeated axis in destination") - - order = [axis for axis in range(x.ndim) if axis not in normalized_source] - for destination_axis, source_axis in sorted( - zip(normalized_destination, normalized_source), - ): - order.insert(destination_axis, source_axis) - return mx.transpose(x, order) - - -def permute_dims(x: Array, axes: tuple[int, ...], /) -> Array: - return mx.transpose(x, axes) - - -def repeat( - x: Array, - repeats: int | Array, - /, - *, - axis: int | None = None, -) -> Array: - if isinstance(repeats, int): - return mx.repeat(x, repeats, axis=axis) - if isinstance(repeats, mx.array) and repeats.ndim == 0: - return mx.repeat(x, int(repeats.item()), axis=axis) - raise NotImplementedError( - "MLX cannot represent the data-dependent output shape produced by " - "a non-scalar repeats array" - ) - - -def reshape( - x: Array, - shape: tuple[int, ...], - /, - *, - copy: bool | None = None, -) -> Array: - result = mx.reshape(x, shape) - return mx.copy(result) if copy is True else result - - -def roll( - x: Array, - shift: int | tuple[int, ...], - /, - *, - axis: int | tuple[int, ...] | None = None, -) -> Array: - return mx.roll(x, shift, axis=axis) - - -def squeeze( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, -) -> Array: - return mx.squeeze(x, axis=axis) - - -def stack(arrays: Sequence[Array], /, *, axis: int = 0) -> Array: - return mx.stack(arrays, axis=axis) - - -def unstack(x: Array, /, *, axis: int = 0) -> tuple[Array, ...]: - return tuple(mx.unstack(x, axis=axis)) - - -def all( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Array: - return mx.all(x, axis=axis, keepdims=keepdims) - - -def any( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Array: - return mx.any(x, axis=axis, keepdims=keepdims) - - -def max( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Array: - return mx.max(x, axis=axis, keepdims=keepdims) - - -def mean( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Array: - return mx.mean(x, axis=axis, keepdims=keepdims) - - -def min( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Array: - return mx.min(x, axis=axis, keepdims=keepdims) - - -def prod( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - dtype: DType | None = None, - keepdims: bool = False, -) -> Array: - if dtype is not None and x.dtype != dtype: - x = mx.astype(x, dtype) - return mx.prod(x, axis=axis, keepdims=keepdims) - - -def std( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - correction: int | float = 0.0, - keepdims: bool = False, -) -> Array: - if not float(correction).is_integer(): - raise ValueError("MLX supports only integral correction values") - return mx.std(x, axis=axis, keepdims=keepdims, ddof=int(correction)) - - -def sum( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - dtype: DType | None = None, - keepdims: bool = False, -) -> Array: - if dtype is not None and x.dtype != dtype: - x = mx.astype(x, dtype) - return mx.sum(x, axis=axis, keepdims=keepdims) - - -def var( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - correction: int | float = 0.0, - keepdims: bool = False, -) -> Array: - if not float(correction).is_integer(): - raise ValueError("MLX supports only integral correction values") - return mx.var(x, axis=axis, keepdims=keepdims, ddof=int(correction)) - - -def argmax( - x: Array, - /, - *, - axis: int | None = None, - keepdims: bool = False, -) -> Array: - return mx.argmax(x, axis=axis, keepdims=keepdims).astype(mx.int32) - - -def argmin( - x: Array, - /, - *, - axis: int | None = None, - keepdims: bool = False, -) -> Array: - return mx.argmin(x, axis=axis, keepdims=keepdims).astype(mx.int32) - - -def count_nonzero( - x: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Array: - return mx.count_nonzero(x, axis=axis, keepdims=keepdims).astype(mx.int32) - - -def _descending_key(x: Array) -> Array: - if x.dtype == mx.bool_: - return mx.logical_not(x) - if mx.issubdtype(x.dtype, mx.unsignedinteger): - return mx.subtract(mx.array(mx.iinfo(x.dtype).max, dtype=x.dtype), x) - return mx.negative(x) - - -def argsort( - x: Array, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, -) -> Array: - # MLX sorting is stable; a stable sort also satisfies stable=False. - key = _descending_key(x) if descending else x - return mx.argsort(key, axis=axis).astype(mx.int32) - - -def sort( - x: Array, - /, - *, - axis: int = -1, - descending: bool = False, - stable: bool = True, -) -> Array: - result = mx.sort(x, axis=axis) - return mx.flip(result, axis=axis) if descending else result - - -def take(x: Array, indices: Array, /, *, axis: int | None = None) -> Array: - return mx.take(x, indices, axis=axis) - - -def take_along_axis( - x: Array, - indices: Array, - /, - *, - axis: int, -) -> Array: - return mx.take_along_axis(x, indices, axis=axis) - - -def clip( - x: Array, - /, - min: int | float | Array | None = None, - max: int | float | Array | None = None, -) -> Array: - if min is None and max is None: - raise ValueError("at least one of min or max must be specified") - return mx.clip(x, min, max) - - -def tril(x: Array, /, *, k: int = 0) -> Array: - return mx.tril(x, k=k) - - -def triu(x: Array, /, *, k: int = 0) -> Array: - return mx.triu(x, k=k) - - -__all__ = [ - "all", - "any", - "arange", - "argmax", - "argmin", - "argsort", - "asarray", - "astype", - "broadcast_arrays", - "broadcast_to", - "clip", - "concat", - "count_nonzero", - "cumulative_prod", - "cumulative_sum", - "diff", - "empty", - "empty_like", - "expand_dims", - "eye", - "flip", - "from_dlpack", - "full", - "full_like", - "linspace", - "matrix_transpose", - "max", - "mean", - "meshgrid", - "min", - "moveaxis", - "ones", - "ones_like", - "permute_dims", - "prod", - "repeat", - "reshape", - "roll", - "sort", - "squeeze", - "stack", - "std", - "sum", - "take", - "take_along_axis", - "tril", - "triu", - "unstack", - "var", - "zeros", - "zeros_like", -] - - -def __dir__() -> list[str]: - return __all__ diff --git a/tests/test_mlx.py b/tests/test_mlx.py index bc2ffd1e..61641d64 100644 --- a/tests/test_mlx.py +++ b/tests/test_mlx.py @@ -7,6 +7,8 @@ import array_api_compat import array_api_compat.mlx as xp +mx.set_default_device(mx.cpu) + def test_namespace_dispatch(): x = mx.arange(3) From 6036035710703456f5621f6af63d57375cac285e Mon Sep 17 00:00:00 2001 From: Declan Healy Date: Sat, 22 Aug 2026 12:16:17 +0100 Subject: [PATCH 04/14] fix MLX cumulative sum and moveaxis aliases --- src/array_api_compat/mlx/_aliases.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/array_api_compat/mlx/_aliases.py b/src/array_api_compat/mlx/_aliases.py index 6c561f84..bf2119e4 100644 --- a/src/array_api_compat/mlx/_aliases.py +++ b/src/array_api_compat/mlx/_aliases.py @@ -301,7 +301,7 @@ def cumulative_sum( normalized_axis = 0 else: normalized_axis = _normalize_axis(axis, x.ndim) - result = mx.cumsum(x, axis=normalized_axis, dtype) + result = mx.cumsum(x, axis=normalized_axis, dtype=dtype) if include_initial: result = _prepend_identity(result, axis=normalized_axis, identity=0) return result @@ -334,3 +334,5 @@ def moveaxis( source: int | tuple[int, ...], destination: int | tuple[int, ...], /, +) -> Array: + return mx.moveaxis(x, source, destination) From 194a054c47249081c9d00d9fe0aefd98b2a3bbf8 Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 12:58:31 +0100 Subject: [PATCH 05/14] Restore complete MLX aliases and fix copy semantics --- src/array_api_compat/mlx/_aliases.py | 357 ++++++++++++++++++++++++++- 1 file changed, 350 insertions(+), 7 deletions(-) diff --git a/src/array_api_compat/mlx/_aliases.py b/src/array_api_compat/mlx/_aliases.py index bf2119e4..5680b358 100644 --- a/src/array_api_compat/mlx/_aliases.py +++ b/src/array_api_compat/mlx/_aliases.py @@ -25,10 +25,7 @@ def _stream(device: Device | None) -> Device | None: def _copy_array(x: Array, *, device: Device | None = None) -> Array: - # ``mlx.core.copy`` is a C++ API but is not currently public in Python. - # ``full_like`` with the input itself as the fill array creates a native - # MLX copy and lets us select the execution device without any host - # conversion. + """Return a native MLX copy, optionally evaluated on *device*.""" return mx.full_like( x, x, @@ -40,7 +37,9 @@ def _copy_array(x: Array, *, device: Device | None = None) -> Array: def _normalize_axis(axis: int, ndim: int) -> int: normalized = axis + ndim if axis < 0 else axis if normalized < 0 or normalized >= ndim: - raise IndexError(f"axis {axis} is out of bounds for an array of dimension {ndim}") + raise IndexError( + f"axis {axis} is out of bounds for an array of dimension {ndim}" + ) return normalized @@ -325,7 +324,9 @@ def flip(x: Array, /, *, axis: int | tuple[int, ...] | None = None) -> Array: def matrix_transpose(x: Array, /) -> Array: if x.ndim < 2: - raise ValueError("matrix_transpose requires an array with at least two dimensions") + raise ValueError( + "matrix_transpose requires an array with at least two dimensions" + ) return mx.swapaxes(x, -1, -2) @@ -335,4 +336,346 @@ def moveaxis( destination: int | tuple[int, ...], /, ) -> Array: - return mx.moveaxis(x, source, destination) + if isinstance(source, int): + if not isinstance(destination, int): + raise ValueError( + "source and destination must have the same number of axes" + ) + return mx.moveaxis(x, source, destination) + + if isinstance(destination, int): + raise ValueError("source and destination must have the same number of axes") + if len(source) != len(destination): + raise ValueError("source and destination must have the same number of axes") + + normalized_source = tuple(_normalize_axis(axis, x.ndim) for axis in source) + normalized_destination = tuple( + _normalize_axis(axis, x.ndim) for axis in destination + ) + if len(set(normalized_source)) != len(normalized_source): + raise ValueError("repeated axis in source") + if len(set(normalized_destination)) != len(normalized_destination): + raise ValueError("repeated axis in destination") + + order = [axis for axis in range(x.ndim) if axis not in normalized_source] + for destination_axis, source_axis in sorted( + zip(normalized_destination, normalized_source), + ): + order.insert(destination_axis, source_axis) + return mx.transpose(x, order) + + +def permute_dims(x: Array, axes: tuple[int, ...], /) -> Array: + return mx.transpose(x, axes) + + +def repeat( + x: Array, + repeats: int | Array, + /, + *, + axis: int | None = None, +) -> Array: + if isinstance(repeats, int): + return mx.repeat(x, repeats, axis=axis) + if isinstance(repeats, mx.array) and repeats.ndim == 0: + return mx.repeat(x, int(repeats.item()), axis=axis) + raise NotImplementedError( + "MLX cannot represent the data-dependent output shape produced by " + "a non-scalar repeats array" + ) + + +def reshape( + x: Array, + shape: tuple[int, ...], + /, + *, + copy: bool | None = None, +) -> Array: + result = mx.reshape(x, shape) + return _copy_array(result) if copy is True else result + + +def roll( + x: Array, + shift: int | tuple[int, ...], + /, + *, + axis: int | tuple[int, ...] | None = None, +) -> Array: + return mx.roll(x, shift, axis=axis) + + +def squeeze( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, +) -> Array: + return mx.squeeze(x, axis=axis) + + +def stack(arrays: Sequence[Array], /, *, axis: int = 0) -> Array: + return mx.stack(arrays, axis=axis) + + +def unstack(x: Array, /, *, axis: int = 0) -> tuple[Array, ...]: + return tuple(mx.unstack(x, axis=axis)) + + +def all( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.all(x, axis=axis, keepdims=keepdims) + + +def any( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.any(x, axis=axis, keepdims=keepdims) + + +def max( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.max(x, axis=axis, keepdims=keepdims) + + +def mean( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.mean(x, axis=axis, keepdims=keepdims) + + +def min( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.min(x, axis=axis, keepdims=keepdims) + + +def prod( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype: DType | None = None, + keepdims: bool = False, +) -> Array: + if dtype is not None and x.dtype != dtype: + x = mx.astype(x, dtype) + return mx.prod(x, axis=axis, keepdims=keepdims) + + +def std( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: int | float = 0.0, + keepdims: bool = False, +) -> Array: + if not float(correction).is_integer(): + raise ValueError("MLX supports only integral correction values") + return mx.std(x, axis=axis, keepdims=keepdims, ddof=int(correction)) + + +def sum( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype: DType | None = None, + keepdims: bool = False, +) -> Array: + if dtype is not None and x.dtype != dtype: + x = mx.astype(x, dtype) + return mx.sum(x, axis=axis, keepdims=keepdims) + + +def var( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: int | float = 0.0, + keepdims: bool = False, +) -> Array: + if not float(correction).is_integer(): + raise ValueError("MLX supports only integral correction values") + return mx.var(x, axis=axis, keepdims=keepdims, ddof=int(correction)) + + +def argmax( + x: Array, + /, + *, + axis: int | None = None, + keepdims: bool = False, +) -> Array: + return mx.argmax(x, axis=axis, keepdims=keepdims).astype(mx.int32) + + +def argmin( + x: Array, + /, + *, + axis: int | None = None, + keepdims: bool = False, +) -> Array: + return mx.argmin(x, axis=axis, keepdims=keepdims).astype(mx.int32) + + +def count_nonzero( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + return mx.count_nonzero(x, axis=axis, keepdims=keepdims).astype(mx.int32) + + +def _descending_key(x: Array) -> Array: + if x.dtype == mx.bool_: + return mx.logical_not(x) + if mx.issubdtype(x.dtype, mx.unsignedinteger): + return mx.subtract(mx.array(mx.iinfo(x.dtype).max, dtype=x.dtype), x) + return mx.negative(x) + + +def argsort( + x: Array, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, +) -> Array: + # MLX sorting is stable; a stable sort also satisfies stable=False. + key = _descending_key(x) if descending else x + return mx.argsort(key, axis=axis).astype(mx.int32) + + +def sort( + x: Array, + /, + *, + axis: int = -1, + descending: bool = False, + stable: bool = True, +) -> Array: + result = mx.sort(x, axis=axis) + return mx.flip(result, axis=axis) if descending else result + + +def take(x: Array, indices: Array, /, *, axis: int | None = None) -> Array: + return mx.take(x, indices, axis=axis) + + +def take_along_axis( + x: Array, + indices: Array, + /, + *, + axis: int, +) -> Array: + return mx.take_along_axis(x, indices, axis=axis) + + +def clip( + x: Array, + /, + min: int | float | Array | None = None, + max: int | float | Array | None = None, +) -> Array: + if min is None and max is None: + raise ValueError("at least one of min or max must be specified") + return mx.clip(x, min, max) + + +def tril(x: Array, /, *, k: int = 0) -> Array: + return mx.tril(x, k=k) + + +def triu(x: Array, /, *, k: int = 0) -> Array: + return mx.triu(x, k=k) + + +__all__ = [ + "all", + "any", + "arange", + "argmax", + "argmin", + "argsort", + "asarray", + "astype", + "broadcast_arrays", + "broadcast_to", + "clip", + "concat", + "count_nonzero", + "cumulative_prod", + "cumulative_sum", + "diff", + "empty", + "empty_like", + "expand_dims", + "eye", + "flip", + "from_dlpack", + "full", + "full_like", + "linspace", + "matrix_transpose", + "max", + "mean", + "meshgrid", + "min", + "moveaxis", + "ones", + "ones_like", + "permute_dims", + "prod", + "repeat", + "reshape", + "roll", + "sort", + "squeeze", + "stack", + "std", + "sum", + "take", + "take_along_axis", + "tril", + "triu", + "unstack", + "var", + "zeros", + "zeros_like", +] + + +def __dir__() -> list[str]: + return __all__ From 8d943b0d36da2af90d4ea17bbd1fc49219b1cfd5 Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 12:59:02 +0100 Subject: [PATCH 06/14] Use native MLX copies for FFT frequency creation --- src/array_api_compat/mlx/fft.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/array_api_compat/mlx/fft.py b/src/array_api_compat/mlx/fft.py index 32608a2b..394d2443 100644 --- a/src/array_api_compat/mlx/fft.py +++ b/src/array_api_compat/mlx/fft.py @@ -8,7 +8,7 @@ import mlx.core as mx from .._internal import clone_module -from ._info import _validate_device +from ._aliases import _copy_array from ._typing import Array, Device, DType __all__ = clone_module("mlx.core.fft", globals()) @@ -160,7 +160,7 @@ def fftfreq( if dtype is not None: result = result.astype(dtype) if device is not None: - result = mx.copy(result, stream=_validate_device(device)) + result = _copy_array(result, device=device) return result @@ -176,7 +176,7 @@ def rfftfreq( if dtype is not None: result = result.astype(dtype) if device is not None: - result = mx.copy(result, stream=_validate_device(device)) + result = _copy_array(result, device=device) return result From 2f1b6f4174656eb7f159c756c909acd10c55d2fa Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 12:59:44 +0100 Subject: [PATCH 07/14] Run MLX decomposition routines on their supported CPU stream --- src/array_api_compat/mlx/linalg.py | 28 +++++++++++++++++----------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/src/array_api_compat/mlx/linalg.py b/src/array_api_compat/mlx/linalg.py index d6b2ee9d..ac678865 100644 --- a/src/array_api_compat/mlx/linalg.py +++ b/src/array_api_compat/mlx/linalg.py @@ -37,7 +37,7 @@ class SVDResult(NamedTuple): def cholesky(x: Array, /, *, upper: bool = False) -> Array: - return mx.linalg.cholesky(x, upper=upper) + return mx.linalg.cholesky(x, upper=upper, stream=mx.cpu) def cross(x1: Array, x2: Array, /, *, axis: int = -1) -> Array: @@ -49,7 +49,7 @@ def diagonal(x: Array, /, *, offset: int = 0) -> Array: def eigh(x: Array, /) -> EighResult: - return EighResult(*mx.linalg.eigh(x)) + return EighResult(*mx.linalg.eigh(x, stream=mx.cpu)) def matrix_norm( @@ -78,7 +78,9 @@ def matrix_rank( singular_values = svdvals(x) largest = mx.max(singular_values, axis=-1, keepdims=True) if rtol is None: - threshold = largest * max(x.shape[-2:]) * mx.finfo(singular_values.dtype).eps + threshold = largest * max(x.shape[-2:]) * mx.finfo( + singular_values.dtype + ).eps else: threshold = largest * mx.asarray(rtol)[..., None] return mx.count_nonzero( @@ -98,9 +100,9 @@ def pinv( rtol: float | Array | None = None, ) -> Array: if rtol is None: - return mx.linalg.pinv(x) + return mx.linalg.pinv(x, stream=mx.cpu) - u, singular_values, vh = mx.linalg.svd(x) + u, singular_values, vh = mx.linalg.svd(x, stream=mx.cpu) largest = mx.max(singular_values, axis=-1, keepdims=True) cutoff = largest * mx.asarray(rtol)[..., None] reciprocal = mx.where( @@ -125,11 +127,11 @@ def qr( raise NotImplementedError( "MLX currently provides reduced QR for tall matrices only" ) - return QRResult(*mx.linalg.qr(x)) + return QRResult(*mx.linalg.qr(x, stream=mx.cpu)) def slogdet(x: Array, /) -> SlogdetResult: - return SlogdetResult(*mx.linalg.slogdet(x)) + return SlogdetResult(*mx.linalg.slogdet(x, stream=mx.cpu)) def svd( @@ -142,11 +144,11 @@ def svd( raise NotImplementedError( "MLX currently provides reduced SVD for rectangular matrices" ) - return SVDResult(*mx.linalg.svd(x)) + return SVDResult(*mx.linalg.svd(x, stream=mx.cpu)) def svdvals(x: Array, /) -> Array: - return mx.linalg.svd(x, compute_uv=False) + return mx.linalg.svd(x, compute_uv=False, stream=mx.cpu) def tensordot( @@ -220,14 +222,18 @@ def vector_norm( order = remaining + axes permuted = mx.transpose(x, order) if order != tuple(range(x.ndim)) else x reduced_size = math.prod(x.shape[index] for index in axes) - reduced_shape = tuple(x.shape[index] for index in remaining) + (reduced_size,) + reduced_shape = tuple(x.shape[index] for index in remaining) + ( + reduced_size, + ) flattened = mx.reshape(permuted, reduced_shape) result = mx.linalg.norm(flattened, ord=ord, axis=-1) if not keepdims: return result - target_shape = [1 if index in axes else x.shape[index] for index in range(x.ndim)] + target_shape = [ + 1 if index in axes else x.shape[index] for index in range(x.ndim) + ] return mx.reshape(result, tuple(target_shape)) From 85e760f621cae83f35d24234e2b651a7601698ba Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 13:02:10 +0100 Subject: [PATCH 08/14] Return concrete MLX devices from namespace inspection --- src/array_api_compat/mlx/_info.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/src/array_api_compat/mlx/_info.py b/src/array_api_compat/mlx/_info.py index ab84e5f6..3fd69f3c 100644 --- a/src/array_api_compat/mlx/_info.py +++ b/src/array_api_compat/mlx/_info.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import mlx.core as mx @@ -17,6 +17,9 @@ if TYPE_CHECKING: from ._typing import DType +_CPU_DEVICE = mx.Device(mx.cpu) +_GPU_DEVICE = mx.Device(mx.gpu) + def _gpu_available() -> bool: """Return whether an MLX GPU backend is available.""" @@ -31,14 +34,22 @@ def _gpu_available() -> bool: return False -def _validate_device(device: Device | None) -> Device: +def _validate_device(device: Any | None) -> Device: + """Normalize MLX device constants to concrete ``mlx.core.Device`` values.""" if device is None: - return mx.default_device() - if not isinstance(device, mx.Device): + selected = mx.default_device() + elif isinstance(device, mx.Device): + selected = device + elif device == mx.cpu: + selected = _CPU_DEVICE + elif device == mx.gpu: + selected = _GPU_DEVICE + else: raise TypeError(f"expected an mlx.core.Device, got {type(device).__name__}") - if device == mx.gpu and not _gpu_available(): + + if selected == _GPU_DEVICE and not _gpu_available(): raise ValueError("the MLX GPU device is not available") - return device + return selected _ALL_DTYPES: dict[str, DType] = { @@ -143,14 +154,14 @@ def dtypes( raise ValueError(f"unsupported kind: {kind!r}") from None # MLX exposes float64 for CPU execution only. - if selected_device == mx.gpu: + if selected_device == _GPU_DEVICE: names = tuple(name for name in names if name != "float64") return {name: _ALL_DTYPES[name] for name in names} def devices(self) -> tuple[Device, ...]: - devices: list[Device] = [mx.cpu] + devices = [_CPU_DEVICE] if _gpu_available(): - devices.append(mx.gpu) + devices.append(_GPU_DEVICE) return tuple(devices) From d22be4ad78874a877c35d3f5f6952535b9e2712f Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 13:22:32 +0100 Subject: [PATCH 09/14] Cancel superseded MLX conformance runs --- .github/workflows/array-api-tests-mlx.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/array-api-tests-mlx.yml b/.github/workflows/array-api-tests-mlx.yml index dea1cd46..aa31f72b 100644 --- a/.github/workflows/array-api-tests-mlx.yml +++ b/.github/workflows/array-api-tests-mlx.yml @@ -9,6 +9,10 @@ on: - main workflow_dispatch: +concurrency: + group: mlx-array-api-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + jobs: tests: runs-on: macos-14 From 8aad0487249fa131e4adc2979b4b63ce9a6e43e7 Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 13:23:38 +0100 Subject: [PATCH 10/14] Add focused MLX Array API semantic overrides --- src/array_api_compat/mlx/_overrides.py | 551 +++++++++++++++++++++++++ 1 file changed, 551 insertions(+) create mode 100644 src/array_api_compat/mlx/_overrides.py diff --git a/src/array_api_compat/mlx/_overrides.py b/src/array_api_compat/mlx/_overrides.py new file mode 100644 index 00000000..35c31131 --- /dev/null +++ b/src/array_api_compat/mlx/_overrides.py @@ -0,0 +1,551 @@ +"""Focused Array API semantic overrides for the MLX compatibility namespace.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from itertools import zip_longest +from typing import Any, NamedTuple + +import mlx.core as mx + +from ._aliases import _copy_array +from ._info import _validate_device +from ._typing import Array, Device, DType + +_DTYPE_TYPE = type(mx.float32) +_INT_BITS = { + mx.int8: 8, + mx.uint8: 8, + mx.int16: 16, + mx.uint16: 16, + mx.int32: 32, + mx.uint32: 32, + mx.int64: 64, + mx.uint64: 64, +} + + +class FInfo(NamedTuple): + bits: int + eps: float + max: float + min: float + smallest_normal: float + dtype: DType + + +class IInfo(NamedTuple): + bits: int + max: int + min: int + dtype: DType + + +class UniqueAllResult(NamedTuple): + values: Array + indices: Array + inverse_indices: Array + counts: Array + + +class UniqueCountsResult(NamedTuple): + values: Array + counts: Array + + +class UniqueInverseResult(NamedTuple): + values: Array + inverse_indices: Array + + +def _dtype_of(value: Any) -> DType: + if isinstance(value, mx.array): + return value.dtype + if isinstance(value, _DTYPE_TYPE): + return value + raise TypeError(f"expected an MLX array or dtype, got {type(value).__name__}") + + +def _as_array(value: Any, *, like: Array | None = None) -> Array: + if isinstance(value, mx.array): + return value + return mx.asarray(value, dtype=None if like is None else like.dtype) + + +def _default_reduction_dtype(dtype: DType) -> DType: + if dtype == mx.bool_ or mx.issubdtype(dtype, mx.signedinteger): + return mx.int32 + if mx.issubdtype(dtype, mx.unsignedinteger): + return mx.uint32 + return dtype + + +def _reduction_size(x: Array, axis: int | tuple[int, ...] | None) -> int: + if axis is None: + return math.prod(x.shape) + axes = (axis,) if isinstance(axis, int) else axis + normalized = tuple(a + x.ndim if a < 0 else a for a in axes) + return math.prod(x.shape[a] for a in normalized) + + +def astype( + x: Array, + dtype: DType, + /, + *, + copy: bool = True, + device: Device | None = None, +) -> Array: + target_device = None if device is None else _validate_device(device) + if x.dtype == dtype: + if not copy and target_device is None: + return x + return _copy_array(x, device=target_device) + result = mx.astype(x, dtype) + return result if target_device is None else _copy_array(result, device=target_device) + + +def broadcast_shapes(*shapes: tuple[int, ...]) -> tuple[int, ...]: + if not shapes: + return () + result: list[int] = [] + for dimensions in zip_longest(*(reversed(shape) for shape in shapes), fillvalue=1): + output = max(dimensions) + if any(dimension not in (1, output) for dimension in dimensions): + raise ValueError(f"shapes {shapes!r} are not broadcastable") + result.append(output) + return tuple(reversed(result)) + + +def broadcast_to(x: Array, /, shape: tuple[int, ...]) -> Array: + return mx.broadcast_to(x, shape) + + +def can_cast(from_: DType | Array, to: DType, /) -> bool: + source = _dtype_of(from_) + try: + return result_type(source, to) == to + except ValueError: + return False + + +def finfo(type_: DType | Array, /) -> FInfo: + dtype = _dtype_of(type_) + real_dtype = mx.float32 if dtype == mx.complex64 else dtype + info = mx.finfo(real_dtype) + bits = 32 if real_dtype == mx.float32 else 64 + smallest_normal = getattr(info, "smallest_normal", getattr(info, "tiny", None)) + return FInfo( + bits=bits, + eps=float(info.eps), + max=float(info.max), + min=float(info.min), + smallest_normal=float(smallest_normal), + dtype=real_dtype, + ) + + +def iinfo(type_: DType | Array, /) -> IInfo: + dtype = _dtype_of(type_) + info = mx.iinfo(dtype) + return IInfo( + bits=_INT_BITS[dtype], + max=int(info.max), + min=int(info.min), + dtype=dtype, + ) + + +def result_type(*arrays_and_dtypes: Any) -> DType: + if not arrays_and_dtypes: + raise ValueError("result_type requires at least one argument") + strong: list[DType] = [] + weak: list[Any] = [] + for value in arrays_and_dtypes: + if isinstance(value, mx.array): + strong.append(value.dtype) + elif isinstance(value, _DTYPE_TYPE): + strong.append(value) + elif isinstance(value, (bool, int, float, complex)): + weak.append(value) + else: + raise TypeError(f"unsupported result_type input {type(value).__name__}") + if strong: + return mx.result_type(*strong) + return mx.result_type(*(mx.asarray(value).dtype for value in weak)) + + +def cumulative_sum( + x: Array, + /, + *, + axis: int | None = None, + dtype: DType | None = None, + include_initial: bool = False, +) -> Array: + if axis is None: + if x.ndim > 1: + raise ValueError("axis must be specified for arrays with more than one dimension") + axis = 0 + output_dtype = _default_reduction_dtype(x.dtype) if dtype is None else dtype + result = mx.cumsum(x, axis=axis, dtype=output_dtype) + if include_initial: + shape = list(result.shape) + shape[axis] = 1 + result = mx.concatenate( + (mx.zeros(tuple(shape), dtype=result.dtype), result), axis=axis + ) + return result + + +def cumulative_prod( + x: Array, + /, + *, + axis: int | None = None, + dtype: DType | None = None, + include_initial: bool = False, +) -> Array: + if axis is None: + if x.ndim > 1: + raise ValueError("axis must be specified for arrays with more than one dimension") + axis = 0 + output_dtype = _default_reduction_dtype(x.dtype) if dtype is None else dtype + result = mx.cumprod(x, axis=axis, dtype=output_dtype) + if include_initial: + shape = list(result.shape) + shape[axis] = 1 + result = mx.concatenate( + (mx.ones(tuple(shape), dtype=result.dtype), result), axis=axis + ) + return result + + +def sum( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype: DType | None = None, + keepdims: bool = False, +) -> Array: + output_dtype = _default_reduction_dtype(x.dtype) if dtype is None else dtype + if x.dtype != output_dtype: + x = mx.astype(x, output_dtype) + return mx.sum(x, axis=axis, keepdims=keepdims) + + +def prod( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype: DType | None = None, + keepdims: bool = False, +) -> Array: + output_dtype = _default_reduction_dtype(x.dtype) if dtype is None else dtype + if x.dtype != output_dtype: + x = mx.astype(x, output_dtype) + return mx.prod(x, axis=axis, keepdims=keepdims) + + +def var( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: int | float = 0.0, + keepdims: bool = False, +) -> Array: + base = mx.var(x, axis=axis, keepdims=keepdims, ddof=0) + count = _reduction_size(x, axis) + numerator = mx.array(count, dtype=base.dtype) + denominator = mx.array(count - correction, dtype=base.dtype) + return base * (numerator / denominator) + + +def std( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: int | float = 0.0, + keepdims: bool = False, +) -> Array: + return mx.sqrt(var(x, axis=axis, correction=correction, keepdims=keepdims)) + + +def diff( + x: Array, + /, + n: int = 1, + axis: int = -1, + prepend: Array | None = None, + append: Array | None = None, +) -> Array: + if n < 0: + raise ValueError("n must be non-negative") + result = x + if prepend is not None: + result = mx.concatenate((prepend, result), axis=axis) + if append is not None: + result = mx.concatenate((result, append), axis=axis) + for _ in range(n): + result = mx.diff(result, axis=axis) + return result + + +def expand_dims( + x: Array, + /, + axis: int | tuple[int, ...], +) -> Array: + axes = (axis,) if isinstance(axis, int) else axis + output_ndim = x.ndim + len(axes) + normalized = tuple(a + output_ndim if a < 0 else a for a in axes) + if any(a < 0 or a >= output_ndim for a in normalized): + raise IndexError("axis is out of bounds") + if len(set(normalized)) != len(normalized): + raise ValueError("repeated axis") + source = iter(x.shape) + shape = tuple(1 if i in normalized else next(source) for i in range(output_ndim)) + return mx.reshape(x, shape) + + +def permute_dims(x: Array, /, axes: tuple[int, ...]) -> Array: + return mx.transpose(x, axes) + + +def reshape( + x: Array, + /, + shape: tuple[int, ...], + *, + copy: bool | None = None, +) -> Array: + result = mx.reshape(x, shape) + return _copy_array(result) if copy is True else result + + +def roll( + x: Array, + /, + shift: int | tuple[int, ...], + *, + axis: int | tuple[int, ...] | None = None, +) -> Array: + return mx.roll(x, shift, axis=axis) + + +def squeeze(x: Array, /, axis: int | tuple[int, ...]) -> Array: + axes = (axis,) if isinstance(axis, int) else axis + normalized = tuple(a + x.ndim if a < 0 else a for a in axes) + if any(a < 0 or a >= x.ndim for a in normalized): + raise IndexError("axis is out of bounds") + if len(set(normalized)) != len(normalized): + raise ValueError("repeated axis") + if any(x.shape[a] != 1 for a in normalized): + raise ValueError("cannot squeeze an axis whose size is not one") + shape = tuple(size for index, size in enumerate(x.shape) if index not in normalized) + return mx.reshape(x, shape) + + +def take_along_axis( + x: Array, + indices: Array, + /, + *, + axis: int = -1, +) -> Array: + return mx.take_along_axis(x, indices, axis=axis) + + +def searchsorted( + x1: Array, + x2: Array | int | float, + /, + *, + side: str = "left", + sorter: Array | None = None, +) -> Array: + sequence = x1 if sorter is None else mx.take(x1, sorter) + values = _as_array(x2, like=x1) + return mx.searchsorted(sequence, values, side=side).astype(mx.int32) + + +def isin( + x1: Array, + x2: Array | int | float | complex | bool, + /, + *, + assume_unique: bool = False, + invert: bool = False, +) -> Array: + del assume_unique + values = _as_array(x2, like=x1).reshape((-1,)) + if values.size == 0: + result = mx.zeros(x1.shape, dtype=mx.bool_) + else: + result = mx.any(x1[..., None] == values, axis=-1) + return mx.logical_not(result) if invert else result + + +def nonzero(x: Array, /) -> tuple[Array, ...]: + raise NotImplementedError("MLX cannot represent data-dependent nonzero shapes") + + +def unique_all(x: Array, /) -> UniqueAllResult: + raise NotImplementedError("MLX cannot represent data-dependent unique shapes") + + +def unique_counts(x: Array, /) -> UniqueCountsResult: + raise NotImplementedError("MLX cannot represent data-dependent unique shapes") + + +def unique_inverse(x: Array, /) -> UniqueInverseResult: + raise NotImplementedError("MLX cannot represent data-dependent unique shapes") + + +def unique_values(x: Array, /) -> Array: + raise NotImplementedError("MLX cannot represent data-dependent unique shapes") + + +def clip( + x: Array, + /, + min: int | float | Array | None = None, + max: int | float | Array | None = None, +) -> Array: + if min is None and max is None: + return _copy_array(x) + return mx.clip(x, min, max) + + +def atan2(x1: Array | int | float, x2: Array | int | float, /) -> Array: + if isinstance(x1, mx.array): + return mx.arctan2(x1, _as_array(x2, like=x1)) + if isinstance(x2, mx.array): + return mx.arctan2(_as_array(x1, like=x2), x2) + return mx.arctan2(mx.asarray(x1), mx.asarray(x2)) + + +def hypot(x1: Array | int | float, x2: Array | int | float, /) -> Array: + if isinstance(x1, mx.array): + a, b = x1, _as_array(x2, like=x1) + elif isinstance(x2, mx.array): + a, b = _as_array(x1, like=x2), x2 + else: + a, b = mx.asarray(x1), mx.asarray(x2) + return mx.sqrt(a * a + b * b) + + +def copysign(x1: Array, x2: Array | int | float, /) -> Array: + raise NotImplementedError("MLX does not yet expose IEEE sign-bit operations") + + +def nextafter(x1: Array, x2: Array | int | float, /) -> Array: + raise NotImplementedError("MLX does not yet expose nextafter") + + +def signbit(x: Array, /) -> Array: + raise NotImplementedError("MLX does not yet expose IEEE sign-bit operations") + + +def expm1(x: Array, /) -> Array: + if x.dtype == mx.complex64: + return mx.exp(x) - mx.ones_like(x) + return mx.expm1(x) + + +def sign(x: Array, /) -> Array: + if x.dtype == mx.complex64: + magnitude = mx.abs(x) + return mx.where(magnitude == 0, mx.zeros_like(x), x / magnitude) + result = mx.sign(x) + if mx.issubdtype(x.dtype, mx.floating): + result = mx.where(mx.isnan(x), x, result) + return result + + +def bitwise_left_shift(x1: Array, x2: Array, /) -> Array: + result = mx.left_shift(x1, x2) + return mx.where(x2 >= _INT_BITS[x1.dtype], mx.zeros_like(result), result) + + +def bitwise_right_shift(x1: Array, x2: Array, /) -> Array: + result = mx.right_shift(x1, x2) + width = _INT_BITS[x1.dtype] + if mx.issubdtype(x1.dtype, mx.signedinteger): + fill = mx.where(x1 < 0, -mx.ones_like(result), mx.zeros_like(result)) + else: + fill = mx.zeros_like(result) + return mx.where(x2 >= width, fill, result) + + +def floor_divide(x1: Array, x2: Array, /) -> Array: + quotient = mx.floor_divide(x1, x2) + if mx.issubdtype(x1.dtype, mx.integer): + remainder = x1 - quotient * x2 + adjust = (remainder != 0) & ((x1 < 0) != (x2 < 0)) + quotient = mx.where(adjust, quotient - 1, quotient) + return quotient + + +def remainder(x1: Array, x2: Array, /) -> Array: + result = x1 - floor_divide(x1, x2) * x2 + positive_zero = mx.zeros_like(result) + negative_zero = -positive_zero + return mx.where( + result == 0, + mx.where(x2 < 0, negative_zero, positive_zero), + result, + ) + + +__all__ = [ + "FInfo", + "IInfo", + "UniqueAllResult", + "UniqueCountsResult", + "UniqueInverseResult", + "astype", + "atan2", + "bitwise_left_shift", + "bitwise_right_shift", + "broadcast_shapes", + "broadcast_to", + "can_cast", + "clip", + "copysign", + "cumulative_prod", + "cumulative_sum", + "diff", + "expand_dims", + "expm1", + "finfo", + "floor_divide", + "hypot", + "iinfo", + "isin", + "nextafter", + "nonzero", + "permute_dims", + "prod", + "remainder", + "reshape", + "result_type", + "roll", + "searchsorted", + "sign", + "signbit", + "squeeze", + "std", + "sum", + "take_along_axis", + "unique_all", + "unique_counts", + "unique_inverse", + "unique_values", + "var", +] From 9066532b7df38283f7b30ae8bbac19cb1d6062c8 Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 13:25:08 +0100 Subject: [PATCH 11/14] Expose MLX Array API semantic overrides --- src/array_api_compat/mlx/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/array_api_compat/mlx/__init__.py b/src/array_api_compat/mlx/__init__.py index b921f3c6..b16b8cf0 100644 --- a/src/array_api_compat/mlx/__init__.py +++ b/src/array_api_compat/mlx/__init__.py @@ -4,8 +4,9 @@ __all__ = clone_module("mlx.core", globals()) -from . import _aliases +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 @@ -18,6 +19,7 @@ __all__ = sorted( set(__all__) | set(_aliases.__all__) + | set(_overrides.__all__) | { "__array_api_version__", "__array_namespace_info__", From 1e22a023b8fa4f6cf8407546c2310f090b4ba5c4 Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 13:25:33 +0100 Subject: [PATCH 12/14] Package MLX semantic overrides --- meson.build | 1 + 1 file changed, 1 insertion(+) diff --git a/meson.build b/meson.build index f3b0b5c2..e4c1ceba 100644 --- a/meson.build +++ b/meson.build @@ -49,6 +49,7 @@ sources_raw = { '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', From 59b311353ac0929630e1982685c72f2a9e7ddabc Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 13:26:07 +0100 Subject: [PATCH 13/14] Report only standardized MLX dtypes --- src/array_api_compat/mlx/_info.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/array_api_compat/mlx/_info.py b/src/array_api_compat/mlx/_info.py index 3fd69f3c..a052b3d5 100644 --- a/src/array_api_compat/mlx/_info.py +++ b/src/array_api_compat/mlx/_info.py @@ -62,7 +62,6 @@ def _validate_device(device: Any | None) -> Device: "uint16": mx.uint16, "uint32": mx.uint32, "uint64": mx.uint64, - "float16": mx.float16, "float32": mx.float32, "float64": mx.float64, "complex64": mx.complex64, @@ -82,7 +81,7 @@ def _validate_device(device: Any | None) -> Device: "uint32", "uint64", ), - "real floating": ("float16", "float32", "float64"), + "real floating": ("float32", "float64"), "complex floating": ("complex64",), "numeric": ( "int8", @@ -93,7 +92,6 @@ def _validate_device(device: Any | None) -> Device: "uint16", "uint32", "uint64", - "float16", "float32", "float64", "complex64", From 215766d7877f8dfa8810d152e399a9dab92208c9 Mon Sep 17 00:00:00 2001 From: declanhealy2 Date: Sat, 22 Aug 2026 13:27:36 +0100 Subject: [PATCH 14/14] Exercise MLX namespace semantic overrides --- tests/test_mlx.py | 61 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 57 insertions(+), 4 deletions(-) diff --git a/tests/test_mlx.py b/tests/test_mlx.py index 61641d64..03693905 100644 --- a/tests/test_mlx.py +++ b/tests/test_mlx.py @@ -23,12 +23,31 @@ def test_namespace_dispatch(): def test_namespace_does_not_patch_array_type(): before_getitem = mx.array.__getitem__ before_namespace = mx.array.__array_namespace__ - # Importing the wrapper must be observational only for the native type. __import__("array_api_compat.mlx") assert mx.array.__getitem__ is before_getitem assert mx.array.__array_namespace__ is before_namespace +def test_namespace_exports_complete_override_surface(): + required = { + "astype", + "broadcast_shapes", + "can_cast", + "cumulative_prod", + "cumulative_sum", + "diff", + "finfo", + "iinfo", + "isin", + "result_type", + "searchsorted", + "sum", + "var", + } + assert required <= set(xp.__all__) + assert all(hasattr(xp, name) for name in required) + + def test_inspection_namespace(): info = xp.__array_namespace_info__() assert info.capabilities() == { @@ -43,6 +62,7 @@ def test_inspection_namespace(): assert defaults["complex floating"] == mx.complex64 assert defaults["integral"] == mx.int32 assert defaults["indexing"] == mx.int32 + assert "float16" not in info.dtypes() assert "complex128" not in info.dtypes() assert info.dtypes(kind="bool") == {"bool": mx.bool_} @@ -62,6 +82,23 @@ def test_creation_signatures_and_copy(): assert isinstance(xp.meshgrid(mx.arange(2), mx.arange(3)), tuple) +def test_dtype_helpers(): + assert xp.broadcast_shapes() == () + assert xp.broadcast_shapes((2, 1), (1, 3)) == (2, 3) + assert xp.can_cast(mx.int8, mx.int16) + assert xp.result_type(mx.int8, mx.int16, 1) == mx.int16 + + float_info = xp.finfo(mx.asarray(1, dtype=mx.float32)) + assert float_info.bits == 32 + assert float_info.dtype == mx.float32 + assert float_info.eps > 0 + + int_info = xp.iinfo(mx.asarray(1, dtype=mx.uint16)) + assert int_info.bits == 16 + assert int_info.dtype == mx.uint16 + assert int_info.min == 0 + + def test_manipulation_wrappers(): x = mx.arange(24).reshape((2, 3, 4)) assert xp.matrix_transpose(x).shape == (2, 4, 3) @@ -69,27 +106,43 @@ def test_manipulation_wrappers(): assert xp.moveaxis(x, (), ()).shape == x.shape assert isinstance(xp.unstack(x), tuple) assert xp.reshape(x, (6, 4), copy=True).shape == (6, 4) + assert xp.expand_dims(mx.arange(3), 1).shape == (3, 1) + assert xp.squeeze(mx.zeros((1, 3, 1)), (0, 2)).shape == (3,) + assert xp.diff(mx.array([0, 1, 3]), n=2).tolist() == [1] def test_reduction_wrappers(): x = mx.arange(6).reshape((2, 3)) assert xp.sum(x, axis=0, dtype=mx.float32).dtype == mx.float32 assert xp.prod(x + 1, axis=1, dtype=mx.float32).dtype == mx.float32 - assert xp.std(x.astype(mx.float32), correction=1).shape == () - assert xp.var(x.astype(mx.float32), correction=1).shape == () + assert xp.sum(mx.array([1], dtype=mx.uint8)).dtype == mx.uint32 + assert xp.prod(mx.array([1], dtype=mx.int8)).dtype == mx.int32 + assert xp.std(x.astype(mx.float32), correction=1.5).shape == () + assert xp.var(x.astype(mx.float32), correction=0.5).shape == () cs = xp.cumulative_sum(mx.array([1, 2, 3]), include_initial=True) cp = xp.cumulative_prod(mx.array([2, 3]), include_initial=True) assert cs.tolist() == [0, 1, 3, 6] assert cp.tolist() == [1, 2, 6] + assert xp.cumulative_sum(mx.array([1], dtype=mx.uint8)).dtype == mx.uint32 -def test_searching_and_sorting_wrappers(): +def test_searching_sorting_and_set_wrappers(): x = mx.array([3, 1, 1, 2]) assert xp.argmax(x).dtype == mx.int32 assert xp.argmin(x).dtype == mx.int32 assert xp.argsort(x, descending=True).tolist() == [0, 3, 1, 2] assert xp.sort(x, descending=True).tolist() == [3, 2, 1, 1] + assert xp.searchsorted(mx.array([1, 3, 5]), 3).item() == 1 + assert xp.isin(mx.array([1, 2, 3]), mx.array([2, 4])).tolist() == [False, True, False] + + +def test_elementwise_overrides(): + assert xp.floor_divide(mx.array(-1), mx.array(2)).item() == -1 + assert xp.remainder(mx.array(-1), mx.array(2)).item() == 1 + assert xp.hypot(mx.array(3.0), 4.0).item() == pytest.approx(5.0) + complex_zero = mx.array(0 + 0j, dtype=mx.complex64) + assert xp.expm1(complex_zero).item() == 0j def test_fft_namespace():