diff --git a/.github/workflows/array-api-tests-mlx.yml b/.github/workflows/array-api-tests-mlx.yml new file mode 100644 index 00000000..aa31f72b --- /dev/null +++ b/.github/workflows/array-api-tests-mlx.yml @@ -0,0 +1,70 @@ +name: MLX Array API Tests + +on: + push: + branches: + - agent/mlx-compat-complete + pull_request: + branches: + - main + workflow_dispatch: + +concurrency: + group: mlx-array-api-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + tests: + runs-on: macos-14 + timeout-minutes: 60 + + steps: + - name: Checkout array-api-compat + uses: actions/checkout@v4 + + - name: Checkout array-api-tests + uses: actions/checkout@v4 + with: + repository: data-apis/array-api-tests + submodules: true + path: array-api-tests + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install mlx pytest + python -m pip install -r array-api-tests/requirements.txt + python -m pip install . + + - name: Run focused MLX wrapper tests + run: python -m pytest tests/test_mlx.py -q + + - name: Run Array API conformance suite + env: + ARRAY_API_TESTS_MODULE: array_api_compat.mlx + ARRAY_API_TESTS_VERSION: "2025.12" + run: | + python - <<'PY' + import mlx.core as mx + import pytest + + mx.set_default_device(mx.cpu) + raise SystemExit( + pytest.main( + [ + "array-api-tests/array_api_tests", + "--max-examples=100", + "--hypothesis-disable-deadline", + "--xfails-file=mlx-xfails.txt", + "--skips-file=mlx-skips.txt", + "-q", + "-rxXfE", + ] + ) + ) + PY diff --git a/meson.build b/meson.build index 1d3accf4..e4c1ceba 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,16 @@ sources_raw = { 'src/array_api_compat/dask/array/linalg.py', ], + 'array_api_compat/mlx': [ + 'src/array_api_compat/mlx/__init__.py', + 'src/array_api_compat/mlx/_aliases.py', + 'src/array_api_compat/mlx/_info.py', + 'src/array_api_compat/mlx/_overrides.py', + 'src/array_api_compat/mlx/_typing.py', + 'src/array_api_compat/mlx/fft.py', + 'src/array_api_compat/mlx/linalg.py', + ], + 'array_api_compat/numpy': [ 'src/array_api_compat/numpy/__init__.py', 'src/array_api_compat/numpy/_aliases.py', 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..9a5db7da --- /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", _base._compat_module_name() + ".mlx"} + + +def array_namespace( + *xs: Any, + api_version: str | None = None, + use_compat: bool | None = None, +) -> ModuleType: + """Return the Array API namespace, including the MLX compat wrapper.""" + mlx_inputs = [x for x in xs if is_mlx_array(x)] + if not mlx_inputs: + return _base.array_namespace( + *xs, + api_version=api_version, + use_compat=use_compat, + ) + + for x in xs: + if isinstance(x, _SCALAR_TYPES): + continue + if not is_mlx_array(x): + raise TypeError("Multiple namespaces for array inputs: MLX and another backend") + + _base._check_api_version(api_version) + if use_compat is False: + import mlx.core as mx + + return mx + + from .. import mlx as mlx_compat + + return mlx_compat + + +get_namespace = array_namespace + + +def is_array_api_obj(x: object) -> bool: + return is_mlx_array(x) or _base.is_array_api_obj(x) + + +def device(x: Any, /) -> Any: + if not is_mlx_array(x): + return _base.device(x) + + import mlx.core as mx + + # MLX arrays use unified memory and do not carry per-array residency. + # The execution default is the only meaningful device value MLX exposes. + return mx.default_device() + + +def to_device( + x: Any, + device: Any, + /, + *, + stream: int | Any | None = None, +) -> Any: + if not is_mlx_array(x): + return _base.to_device(x, device, stream=stream) + if stream is not None: + raise NotImplementedError("MLX does not expose Array API stream handles") + + import mlx.core as mx + + if not isinstance(device, mx.Device): + raise TypeError(f"expected an mlx.core.Device, got {type(device).__name__}") + return mx.full_like(x, x, dtype=x.dtype, stream=device) + + +def is_lazy_array(x: object) -> bool: + if is_mlx_array(x): + return True + return _base.is_lazy_array(x) + + +__all__ = [ + "array_namespace", + "device", + "get_namespace", + "is_array_api_obj", + "is_lazy_array", + "is_mlx_array", + "is_mlx_namespace", + "to_device", +] + + +def __dir__() -> list[str]: + return __all__ diff --git a/src/array_api_compat/mlx/__init__.py b/src/array_api_compat/mlx/__init__.py new file mode 100644 index 00000000..b16b8cf0 --- /dev/null +++ b/src/array_api_compat/mlx/__init__.py @@ -0,0 +1,33 @@ +from typing import Final + +from .._internal import clone_module + +__all__ = clone_module("mlx.core", globals()) + +from . import _aliases, _overrides +from ._aliases import * # type: ignore[assignment,no-redef] # noqa: F403 +from ._overrides import * # type: ignore[assignment,no-redef] # noqa: F403 +from ._info import __array_namespace_info__ + +# Import the compatibility submodules explicitly so they replace the native +# ``mlx.core.fft`` and ``mlx.core.linalg`` objects cloned above. +fft = __import__(__spec__.parent + ".fft", fromlist=["fft"]) +linalg = __import__(__spec__.parent + ".linalg", fromlist=["linalg"]) + +__array_api_version__: Final = "2025.12" + +__all__ = sorted( + set(__all__) + | set(_aliases.__all__) + | set(_overrides.__all__) + | { + "__array_api_version__", + "__array_namespace_info__", + "fft", + "linalg", + } +) + + +def __dir__() -> list[str]: + return __all__ diff --git a/src/array_api_compat/mlx/_aliases.py b/src/array_api_compat/mlx/_aliases.py new file mode 100644 index 00000000..5680b358 --- /dev/null +++ b/src/array_api_compat/mlx/_aliases.py @@ -0,0 +1,681 @@ +"""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 _copy_array(x: Array, *, device: Device | None = None) -> Array: + """Return a native MLX copy, optionally evaluated on *device*.""" + 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: + 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") + + 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 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( + 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 _copy_array(result, device=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 _copy_array(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 _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__ diff --git a/src/array_api_compat/mlx/_info.py b/src/array_api_compat/mlx/_info.py new file mode 100644 index 00000000..a052b3d5 --- /dev/null +++ b/src/array_api_compat/mlx/_info.py @@ -0,0 +1,170 @@ +"""Array API inspection support for MLX.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import mlx.core as mx + +from ..common._typing import ( + Capabilities, + DTypeKind, + DTypesAny, + DefaultDTypes, +) +from ._typing import Device + +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.""" + 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: Any | None) -> Device: + """Normalize MLX device constants to concrete ``mlx.core.Device`` values.""" + if device is None: + 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 selected == _GPU_DEVICE and not _gpu_available(): + raise ValueError("the MLX GPU device is not available") + return selected + + +_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, + "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": ("float32", "float64"), + "complex floating": ("complex64",), + "numeric": ( + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "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 == _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 = [_CPU_DEVICE] + if _gpu_available(): + devices.append(_GPU_DEVICE) + return tuple(devices) + + +__all__ = ["__array_namespace_info__"] + + +def __dir__() -> list[str]: + return __all__ 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", +] 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..394d2443 --- /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 ._aliases import _copy_array +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 = _copy_array(result, 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 = _copy_array(result, 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..ac678865 --- /dev/null +++ b/src/array_api_compat/mlx/linalg.py @@ -0,0 +1,269 @@ +"""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, stream=mx.cpu) + + +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, stream=mx.cpu)) + + +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, stream=mx.cpu) + + 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( + 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, stream=mx.cpu)) + + +def slogdet(x: Array, /) -> SlogdetResult: + return SlogdetResult(*mx.linalg.slogdet(x, stream=mx.cpu)) + + +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, stream=mx.cpu)) + + +def svdvals(x: Array, /) -> Array: + return mx.linalg.svd(x, compute_uv=False, stream=mx.cpu) + + +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..03693905 --- /dev/null +++ b/tests/test_mlx.py @@ -0,0 +1,166 @@ +import inspect + +import pytest + +mx = pytest.importorskip("mlx.core") + +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) + 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__ + __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() == { + "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 "float16" not in info.dtypes() + 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_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) + 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) + 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.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_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(): + 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 == ()