diff --git a/.github/workflows/mlx-tests.yml b/.github/workflows/mlx-tests.yml new file mode 100644 index 00000000..2996f9cf --- /dev/null +++ b/.github/workflows/mlx-tests.yml @@ -0,0 +1,24 @@ +name: MLX Tests + +on: [push, pull_request] + +jobs: + mlx-tests: + runs-on: macos-14 # Apple Silicon runner + steps: + - uses: actions/checkout@v7.0.0 + + - name: Set up Python + uses: actions/setup-python@v6.2.0 + with: + python-version: '3.13' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install pytest mlx numpy array-api-strict ndonnx + + - name: Run MLX Compat Tests + run: | + export PYTHONPATH="${GITHUB_WORKSPACE}/src" + pytest tests/test_mlx.py tests/test_common.py tests/test_isdtype.py tests/test_copies_or_views.py -k mlx -v diff --git a/meson.build b/meson.build index 1d3accf4..972cd74a 100644 --- a/meson.build +++ b/meson.build @@ -61,6 +61,15 @@ sources_raw = { 'src/array_api_compat/torch/fft.py', 'src/array_api_compat/torch/linalg.py', ], + + 'array_api_compat/mlx': [ + 'src/array_api_compat/mlx/__init__.py', + 'src/array_api_compat/mlx/_aliases.py', + 'src/array_api_compat/mlx/_info.py', + 'src/array_api_compat/mlx/_typing.py', + 'src/array_api_compat/mlx/fft.py', + 'src/array_api_compat/mlx/linalg.py', + ], } sources = {} diff --git a/pyproject.toml b/pyproject.toml index 4ca2f429..da7f961a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dev = [ "pytest", "torch", "sparse>=0.15.1", + "mlx; platform_system == 'Darwin' and platform_machine == 'arm64'" ] diff --git a/src/array_api_compat/common/_helpers.py b/src/array_api_compat/common/_helpers.py index d5342658..cc216f3a 100644 --- a/src/array_api_compat/common/_helpers.py +++ b/src/array_api_compat/common/_helpers.py @@ -116,7 +116,7 @@ def is_numpy_array(x: object) -> TypeIs[npt.NDArray[Any]]: # TODO: Should we reject ndarray subclasses? cls = cast(Hashable, type(x)) return ( - _issubclass_fast(cls, "numpy", "ndarray") + _issubclass_fast(cls, "numpy", "ndarray") or _issubclass_fast(cls, "numpy", "generic") ) and not _is_jax_zero_gradient_array(x) @@ -273,6 +273,29 @@ def is_pydata_sparse_array(x: object) -> TypeIs[sparse.SparseArray]: return _issubclass_fast(cls, "sparse", "SparseArray") +def is_mlx_array(x: object) -> bool: + """ + Return True if `x` is an MLX array. + + This function does not import MLX if it has not already been imported + and is therefore cheap to use. + + See Also + -------- + + array_namespace + is_array_api_obj + is_numpy_array + is_cupy_array + is_torch_array + is_dask_array + is_jax_array + is_pydata_sparse_array + """ + cls = cast(Hashable, type(x)) + return _issubclass_fast(cls, "mlx.core", "array") + + def is_array_api_obj(x: object) -> TypeGuard[_ArrayApiObj]: """ Return True if `x` is an array API compatible array object. @@ -313,6 +336,7 @@ def _is_array_api_cls(cls: type) -> bool: # TODO: drop support for jax<0.4.32 which didn't have __array_namespace__ or _issubclass_fast(cls, "jax", "Array") or _issubclass_fast(cls, "jax.core", "Tracer") # see is_jax_array for limitations + or _issubclass_fast(cls, "mlx.core", "array") ) @@ -488,6 +512,27 @@ def is_array_api_strict_namespace(xp: Namespace) -> bool: return xp.__name__ == "array_api_strict" +@lru_cache(100) +def is_mlx_namespace(xp: Namespace) -> bool: + """ + Returns True if `xp` is an MLX namespace. + + This includes both ``mlx.core`` itself and the version wrapped by array-api-compat. + + See Also + -------- + + array_namespace + is_numpy_namespace + is_cupy_namespace + is_torch_namespace + is_dask_namespace + is_jax_namespace + is_pydata_sparse_namespace + """ + return xp.__name__ in {"mlx.core", _compat_module_name() + ".mlx"} + + def _check_api_version(api_version: str | None) -> None: if api_version in _API_VERSIONS_OLD: warnings.warn( @@ -516,7 +561,7 @@ def _cls_to_namespace( cls_ = cast(Hashable, cls) # Make mypy happy if ( - _issubclass_fast(cls_, "numpy", "ndarray") + _issubclass_fast(cls_, "numpy", "ndarray") or _issubclass_fast(cls_, "numpy", "generic") ): if use_compat is True: @@ -559,6 +604,14 @@ def _cls_to_namespace( import dask.array as xp # type: ignore[no-redef] return xp, None + if _issubclass_fast(cls_, "mlx.core", "array"): + if _use_compat: + _check_api_version(api_version) + from .. import mlx as xp # type: ignore[no-redef] + else: + import mlx.core as xp # type: ignore[no-redef] + return xp, None + # Backwards compatibility for jax<0.4.32 if _issubclass_fast(cls_, "jax", "Array"): return _jax_namespace(api_version, use_compat), None @@ -731,12 +784,45 @@ def __repr__(self) -> Literal["DASK_DEVICE"]: _DASK_DEVICE = _dask_device() -# device() is not on numpy.ndarray or dask.array and to_device() is not on numpy.ndarray -# or cupy.ndarray. They are not included in array objects of this library -# because this library just reuses the respective ndarray classes without -# wrapping or subclassing them. These helper functions can be used instead of -# the wrapper functions for libraries that need to support both NumPy/CuPy and -# other libraries that use devices. +_mlx_device_map = {} + + +def _mlx_cleanup(ref_id): + _mlx_device_map.pop(ref_id, None) + + +def _get_mlx_device(x): + import mlx.core as mx + return _mlx_device_map.get(id(x), mx.default_device()) + + +def _set_mlx_device(x, device): + import weakref + ref_id = id(x) + _mlx_device_map[ref_id] = device + try: + # Register callback to pop ref_id when x is garbage collected. + # This prevents memory leaks without keeping a strong reference to x. + weakref.finalize(x, _mlx_cleanup, ref_id) + except TypeError: + # If x doesn't support weak references (though mlx arrays do) + pass + + +def _normalize_mlx_device(device): + import mlx.core as mx + if isinstance(device, str): + if device == "cpu": + return mx.Device(mx.DeviceType.cpu) + elif device == "gpu": + return mx.Device(mx.DeviceType.gpu) + else: + raise ValueError(f"Unsupported device: {device!r}") + if not isinstance(device, mx.Device): + raise TypeError(f"Unsupported device type: {type(device)}") + return device + + def device(x: _ArrayApiObj, /) -> Device: """ Hardware device the array data resides on. @@ -802,6 +888,8 @@ def device(x: _ArrayApiObj, /) -> Device: return "cpu" # Return the device of the constituent array return device(inner) # pyright: ignore + elif is_mlx_array(x): + return _get_mlx_device(x) return x.device # type: ignore # pyright: ignore @@ -927,6 +1015,16 @@ def to_device(x: Array, device: Device, /, *, stream: int | Any | None = None) - if not hasattr(x, "to_device"): return x return x.to_device(device, stream=stream) + elif is_mlx_array(x): + if stream is not None: + raise ValueError("The stream argument to to_device() is not supported for MLX") + dev = _normalize_mlx_device(device) + import mlx.core as mx + # MLX unified memory, so we just construct a new array reference + # and store its mapped device in the id->device registry (_mlx_device_map). + y = mx.array(x) + _set_mlx_device(y, dev) + return y elif is_pydata_sparse_array(x) and device == _device(x): # Perform trivial check to return the same array if # device is same instead of err-ing. @@ -1079,6 +1177,8 @@ def is_lazy_array(x: object) -> TypeGuard[_ArrayApiObj]: "is_dask_namespace", "is_jax_array", "is_jax_namespace", + "is_mlx_array", + "is_mlx_namespace", "is_numpy_array", "is_numpy_namespace", "is_torch_array", diff --git a/src/array_api_compat/mlx/__init__.py b/src/array_api_compat/mlx/__init__.py new file mode 100644 index 00000000..33f431f2 --- /dev/null +++ b/src/array_api_compat/mlx/__init__.py @@ -0,0 +1,77 @@ +from typing import Final + +from .._internal import clone_module + + +__all__ = clone_module("mlx.core", globals()) + +# These imports may overwrite names from the import * above. +from . import _aliases +from ._aliases import * # type: ignore[assignment,no-redef] # noqa: F403 +from ._info import __array_namespace_info__ # noqa: F401 + +# Don't know why, but we have to do an absolute import to import linalg. If we +# instead do +# +# from . import linalg +# +# It doesn't overwrite np.linalg from above. The import is generated +# dynamically so that the library can be vendored. +__import__(__spec__.parent + ".linalg") + +__import__(__spec__.parent + ".fft") + +from .linalg import matrix_transpose, vecdot # type: ignore[no-redef] # noqa: F401 + +__array_api_version__: Final = "2025.12" + +__all__ = sorted( + set(__all__) + | set(_aliases.__all__) + | {"__array_api_version__", "__array_namespace_info__", "linalg", "fft"} +) + +def __dir__() -> list[str]: + return __all__ + + +# Monkeypatch mlx.core.array to support boolean indexing __getitem__ +# MLX doesn't natively support boolean indexing __getitem__ (data-dependent shapes), +# so we provide a NumPy-based fallback. +import mlx.core as mx +import numpy as np +import builtins + +_old_getitem = mx.array.__getitem__ + +def _is_bool_array(x): + return isinstance(x, mx.array) and x.dtype == mx.bool_ + +def _has_bool_array(item): + if _is_bool_array(item): + return True + if isinstance(item, tuple): + return builtins.any(_has_bool_array(i) for i in item) + if isinstance(item, list): + return builtins.any(_has_bool_array(i) for i in item) + return False + +def _to_numpy_indices(item): + if _is_bool_array(item): + return np.array(item) + if isinstance(item, tuple): + return tuple(_to_numpy_indices(i) for i in item) + if isinstance(item, list): + return [_to_numpy_indices(i) for i in item] + return item + +def _new_getitem(self, item): + if _has_bool_array(item): + # Convert self and any boolean masks to numpy + self_np = np.array(self) + item_np = _to_numpy_indices(item) + res_np = self_np[item_np] + return mx.array(res_np) + return _old_getitem(self, item) + +mx.array.__getitem__ = _new_getitem diff --git a/src/array_api_compat/mlx/_aliases.py b/src/array_api_compat/mlx/_aliases.py new file mode 100644 index 00000000..57a854c7 --- /dev/null +++ b/src/array_api_compat/mlx/_aliases.py @@ -0,0 +1,383 @@ +# pyright: reportPrivateUsage=false +from __future__ import annotations + +from builtins import bool as py_bool +from typing import Any + +import mlx.core as mx + +from .._internal import get_xp +from ..common import _aliases +from ..common._helpers import _check_device +from ..common._typing import NestedSequence, SupportsBufferProtocol +from ._typing import Array, Device, DType + +bool = mx.bool_ + +# --- Renamed trig functions (MLX uses arc* like NumPy) --- +acos = mx.arccos +acosh = mx.arccosh +asin = mx.arcsin +asinh = mx.arcsinh +atan = mx.arctan +atan2 = mx.arctan2 +atanh = mx.arctanh + +# --- Bitwise renames (same as NumPy) --- +bitwise_left_shift = mx.left_shift +bitwise_right_shift = mx.right_shift +bitwise_invert = mx.bitwise_invert + +# --- concat -> concatenate --- +concat = mx.concatenate + +# --- pow --- +pow = mx.power + +# --- Creation functions: delegate to common (adds device kwarg handling) --- +# arange/eye/linspace need MLX-specific overrides due to kwarg name differences. +empty = get_xp(mx)(_aliases.empty) +empty_like = get_xp(mx)(_aliases.empty_like) +full = get_xp(mx)(_aliases.full) +full_like = get_xp(mx)(_aliases.full_like) +ones = get_xp(mx)(_aliases.ones) +ones_like = get_xp(mx)(_aliases.ones_like) +zeros = get_xp(mx)(_aliases.zeros) +zeros_like = get_xp(mx)(_aliases.zeros_like) + + +def arange( + start: float, + /, + stop: float | None = None, + step: float = 1, + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: Any, +) -> Array: + """Array API wrapper: MLX arange doesn't accept stop=None as keyword.""" + _check_device(mx, device) + if stop is None: + # MLX arange(stop) form, step ignored when only stop given + return mx.arange(start, dtype=dtype) + return mx.arange(start, stop, step, dtype=dtype) + + +def eye( + n_rows: int, + n_cols: int | None = None, + /, + *, + k: int = 0, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: Any, +) -> Array: + """Array API wrapper: MLX eye uses 'm' not 'M'.""" + _check_device(mx, device) + return mx.eye(n_rows, m=n_cols, k=k, dtype=dtype) + + +def linspace( + start: float, + stop: float, + /, + num: int = 50, + *, + dtype: DType | None = None, + device: Device | None = None, + endpoint: py_bool = True, + **kwargs: Any, +) -> Array: + """Array API wrapper: MLX linspace has no endpoint kwarg.""" + _check_device(mx, device) + if not endpoint: + # Emulate endpoint=False: generate num+1 points, drop the last + return mx.linspace(start, stop, num + 1, dtype=dtype)[:-1] + return mx.linspace(start, stop, num, dtype=dtype) + +# --- Unique (split from mx.unique like NumPy) --- +UniqueAllResult = get_xp(mx)(_aliases.UniqueAllResult) +UniqueCountsResult = get_xp(mx)(_aliases.UniqueCountsResult) +UniqueInverseResult = get_xp(mx)(_aliases.UniqueInverseResult) +unique_all = get_xp(mx)(_aliases.unique_all) +unique_counts = get_xp(mx)(_aliases.unique_counts) +unique_inverse = get_xp(mx)(_aliases.unique_inverse) +unique_values = get_xp(mx)(_aliases.unique_values) + +# --- std/var: MLX requires ddof as int, not float --- +def std( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: float = 0.0, + keepdims: py_bool = False, + **kwargs: Any, +) -> Array: + """Array API wrapper: MLX std requires ddof as int.""" + return mx.std(x, axis=axis, ddof=int(correction), keepdims=keepdims) + + +def var( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: float = 0.0, + keepdims: py_bool = False, + **kwargs: Any, +) -> Array: + """Array API wrapper: MLX var requires ddof as int.""" + return mx.var(x, axis=axis, ddof=int(correction), keepdims=keepdims) + +# --- cumulative_sum/prod: rename from cumsum/cumprod + include_initial --- +cumulative_sum = get_xp(mx)(_aliases.cumulative_sum) +cumulative_prod = get_xp(mx)(_aliases.cumulative_prod) + +def clip( + x: Array, + /, + min: float | Array | None = None, + max: float | Array | None = None, + **kwargs: Any, +) -> Array: + """ + Array API compatibility wrapper for clip(). + """ + if min is None and max is None: + return mx.array(x) + if min is None: + return mx.minimum(x, mx.array(max, dtype=x.dtype)) + if max is None: + return mx.maximum(x, mx.array(min, dtype=x.dtype)) + return mx.clip(x, mx.array(min, dtype=x.dtype), mx.array(max, dtype=x.dtype)) + + +# --- permute_dims: MLX has this natively (unlike NumPy which has transpose) --- +permute_dims = mx.permute_dims + + +# --- reshape: MLX uses 'shape' not 'newshape', but copy kwarg unsupported --- +def reshape( + x: Array, + /, + shape: tuple[int, ...], + *, + copy: py_bool | None = None, +) -> Array: + """ + Array API compatibility wrapper for reshape(). + + See the corresponding documentation in MLX and/or the array API + specification for more details. + """ + if copy is True: + x = mx.array(x) + return mx.reshape(x, shape) + + +# --- argsort/sort: add descending + stable kwargs MLX doesn't have --- +def argsort( + x: Array, + /, + *, + axis: int = -1, + descending: py_bool = False, + stable: py_bool = True, +) -> Array: + """ + Array API compatibility wrapper for argsort(). + + See the corresponding documentation in MLX and/or the array API + specification for more details. + """ + # MLX argsort has no stable/descending; stable sort is the default. + # descending is emulated via flip. + if not descending: + return mx.argsort(x, axis=axis) + res = mx.flip( + mx.argsort(mx.flip(x, axis=axis), axis=axis), + axis=axis, + ) + normalised_axis = axis if axis >= 0 else x.ndim + axis + max_i = x.shape[normalised_axis] - 1 + return max_i - res + + +def sort( + x: Array, + /, + *, + axis: int = -1, + descending: py_bool = False, + stable: py_bool = True, +) -> Array: + """ + Array API compatibility wrapper for sort(). + + See the corresponding documentation in MLX and/or the array API + specification for more details. + """ + # MLX sort has no stable/descending kwargs + res = mx.sort(x, axis=axis) + if descending: + res = mx.flip(res, axis=axis) + return res + + +# --- nonzero: must error on 0-d --- +nonzero = get_xp(mx)(_aliases.nonzero) + +# --- linalg helpers --- +matmul = get_xp(mx)(_aliases.matmul) +matrix_transpose = get_xp(mx)(_aliases.matrix_transpose) +tensordot = get_xp(mx)(_aliases.tensordot) +sign = get_xp(mx)(_aliases.sign) + +# --- vecdot: MLX has no native vecdot --- +vecdot = get_xp(mx)(_aliases.vecdot) + +# --- isdtype: MLX has DtypeCategory but no isdtype function --- +isdtype = get_xp(mx)(_aliases.isdtype) + +# --- unstack: MLX has no native unstack --- +unstack = get_xp(mx)(_aliases.unstack) + +# --- finfo --- +finfo = get_xp(mx)(_aliases.finfo) + + +class iinfo_result: + def __init__(self, min: int, max: int, bits: int, dtype: DType): + self.min = min + self.max = max + self.bits = bits + self.dtype = dtype + + def __repr__(self) -> str: + return f"iinfo(min={self.min}, max={self.max}, dtype={self.dtype})" + + +def iinfo(type_: DType | Array, /) -> Any: + """ + Array API compatibility wrapper for iinfo(). + + MLX does not have a native iinfo(); we provide a minimal shim that returns + an object with .min, .max, .bits, and .dtype attributes for integer dtypes. + + See the corresponding documentation in the array API specification for more + details. + """ + # Resolve dtype from array if needed + if isinstance(type_, mx.array): + dtype = type_.dtype + else: + dtype = type_ + + # MLX has no iinfo; hand-coded table for integer dtypes it supports + _info = { + mx.int8: (-(2**7), 2**7 - 1, 8, mx.int8), + mx.int16: (-(2**15), 2**15 - 1, 16, mx.int16), + mx.int32: (-(2**31), 2**31 - 1, 32, mx.int32), + mx.int64: (-(2**63), 2**63 - 1, 64, mx.int64), + mx.uint8: (0, 2**8 - 1, 8, mx.uint8), + mx.uint16: (0, 2**16 - 1, 16, mx.uint16), + mx.uint32: (0, 2**32 - 1, 32, mx.uint32), + mx.uint64: (0, 2**64 - 1, 64, mx.uint64), + } + if dtype not in _info: + raise TypeError(f"iinfo is not supported for dtype {dtype}") + + mn, mx_, bits, dt = _info[dtype] + return iinfo_result(mn, mx_, bits, dt) + + +# --- asarray: MLX has no asarray; wrap mx.array with spec-compatible signature --- +def asarray( + obj: Array | complex | NestedSequence[complex] | SupportsBufferProtocol, + /, + *, + dtype: DType | None = None, + device: Device | None = None, + copy: py_bool | None = None, + **kwargs: Any, +) -> Array: + """ + Array API compatibility wrapper for asarray(). + """ + _check_device(mx, device) + if isinstance(obj, mx.array): + if dtype is None: + dtype = obj.dtype + if copy is False: + if dtype != obj.dtype: + raise ValueError("Cannot perform copy=False with a different dtype") + return obj + elif copy is None or copy is True: + if copy is True: + return mx.array(obj, dtype=dtype) + else: + if dtype == obj.dtype: + return obj + return mx.array(obj, dtype=dtype) + return mx.array(obj, dtype=dtype) + + +# --- astype: MLX astype has no copy kwarg --- +def astype( + x: Array, + dtype: DType, + /, + *, + copy: py_bool = True, + device: Device | None = None, +) -> Array: + """ + Array API compatibility wrapper for astype(). + + See the corresponding documentation in MLX and/or the array API + specification for more details. + """ + # MLX astype has no copy param; always returns new array + return x.astype(dtype) + + +# --- take_along_axis: axis defaults to -1 in spec, required in MLX --- +def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array: + """ + Array API compatibility wrapper for take_along_axis(). + + See the corresponding documentation in MLX and/or the array API + specification for more details. + """ + return mx.take_along_axis(x, indices, axis=axis) + + +__all__ = _aliases.__all__ + [ + "asarray", + "astype", + "acos", + "acosh", + "asin", + "asinh", + "atan", + "atan2", + "atanh", + "bitwise_left_shift", + "bitwise_invert", + "bitwise_right_shift", + "bool", + "concat", + "pow", + "iinfo", + "reshape", + "argsort", + "sort", + "take_along_axis", +] + + +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..13ad788b --- /dev/null +++ b/src/array_api_compat/mlx/_info.py @@ -0,0 +1,164 @@ +""" +Array API Inspection namespace for MLX. + +See https://data-apis.org/array-api/latest/API_specification/inspection.html +for more details. +""" +from __future__ import annotations + +import mlx.core as mx + +from ..common._typing import DefaultDTypes +from ._typing import Device, DType + + +class __array_namespace_info__: + """ + Get the array API inspection namespace for MLX. + + The array API inspection namespace defines the following functions: + + - capabilities() + - default_device() + - default_dtypes() + - dtypes() + - devices() + + See + https://data-apis.org/array-api/latest/API_specification/inspection.html + for more details. + """ + + __module__ = "mlx.core" + + def capabilities(self) -> dict: + """ + Return a dictionary of array API library capabilities. + + For MLX: + - "boolean indexing": True + - "data-dependent shapes": False (MLX uses lazy evaluation with static shapes) + - "max dimensions": 16 (MLX practical limit) + + See + https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html + for more details. + """ + return { + "boolean indexing": True, + # MLX uses lazy/static shapes; nonzero() output size is + # not known at trace time, so data-dependent shapes are not supported. + "data-dependent shapes": False, + "max dimensions": 16, + } + + def default_device(self) -> Device: + """ + The default device used for new MLX arrays. + + Returns the MLX default device (typically the GPU on Apple Silicon). + """ + return mx.default_device() + + def default_dtypes( + self, + *, + device: Device | None = None, + ) -> DefaultDTypes: + """ + The default data types used for new MLX arrays. + + MLX defaults: + - "real floating": float32 (MLX has no float64 by default) + - "complex floating": complex64 + - "integral": int32 + - "indexing": int32 + """ + # MLX doesn't validate devices per call; unified memory model. + return { + "real floating": mx.float32, + "complex floating": mx.complex64, + "integral": mx.int32, + "indexing": mx.int32, + } + + def dtypes( + self, + *, + device: Device | None = None, + kind: str | tuple[str, ...] | None = None, + ) -> dict[str, DType]: + """ + The array API data types supported by MLX. + + Note: MLX does not support float64 or complex128. + + Parameters + ---------- + device : optional + Ignored (MLX uses unified memory). + kind : str or tuple of str, optional + Filter by kind: 'bool', 'signed integer', 'unsigned integer', + 'integral', 'real floating', 'complex floating', 'numeric'. + """ + _bool = {"bool": mx.bool_} + _signed = { + "int8": mx.int8, + "int16": mx.int16, + "int32": mx.int32, + "int64": mx.int64, + } + _unsigned = { + "uint8": mx.uint8, + "uint16": mx.uint16, + "uint32": mx.uint32, + "uint64": mx.uint64, + } + # MLX has no float64/complex128 + _real_float = { + "float16": mx.float16, + "bfloat16": mx.bfloat16, + "float32": mx.float32, + } + _complex_float = { + "complex64": mx.complex64, + } + + if kind is None: + return {**_bool, **_signed, **_unsigned, **_real_float, **_complex_float} + if kind == "bool": + return _bool + if kind == "signed integer": + return _signed + if kind == "unsigned integer": + return _unsigned + if kind == "integral": + return {**_signed, **_unsigned} + if kind == "real floating": + return _real_float + if kind == "complex floating": + return _complex_float + if kind == "numeric": + return {**_signed, **_unsigned, **_real_float, **_complex_float} + if isinstance(kind, tuple): + res: dict[str, DType] = {} + for k in kind: + res.update(self.dtypes(kind=k)) + return res + raise ValueError(f"unsupported kind: {kind!r}") + + def devices(self) -> list[Device]: + """ + The devices supported by MLX. + + Returns CPU and GPU devices available on the current machine. + """ + # Return both available MLX device types + return [mx.Device(mx.DeviceType.cpu), mx.Device(mx.DeviceType.gpu)] + + +__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..bd44c310 --- /dev/null +++ b/src/array_api_compat/mlx/_typing.py @@ -0,0 +1,3 @@ +from mlx.core import array as Array, Dtype as DType, Device + +__all__ = ["Array", "DType", "Device"] diff --git a/src/array_api_compat/mlx/fft.py b/src/array_api_compat/mlx/fft.py new file mode 100644 index 00000000..c6dfdcd4 --- /dev/null +++ b/src/array_api_compat/mlx/fft.py @@ -0,0 +1,28 @@ +import mlx.core as mx + +from .._internal import clone_module, get_xp +from ..common import _fft + +__all__ = clone_module("mlx.core.fft", globals()) + +# Wrap all FFT functions via common helpers (handles float32->complex64 upcasting) +fft = get_xp(mx)(_fft.fft) +ifft = get_xp(mx)(_fft.ifft) +fftn = get_xp(mx)(_fft.fftn) +ifftn = get_xp(mx)(_fft.ifftn) +rfft = get_xp(mx)(_fft.rfft) +irfft = get_xp(mx)(_fft.irfft) +rfftn = get_xp(mx)(_fft.rfftn) +irfftn = get_xp(mx)(_fft.irfftn) +hfft = get_xp(mx)(_fft.hfft) +ihfft = get_xp(mx)(_fft.ihfft) +fftfreq = get_xp(mx)(_fft.fftfreq) +rfftfreq = get_xp(mx)(_fft.rfftfreq) +fftshift = get_xp(mx)(_fft.fftshift) +ifftshift = get_xp(mx)(_fft.ifftshift) + +__all__ = sorted(set(__all__) | set(_fft.__all__)) + + +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..edf30011 --- /dev/null +++ b/src/array_api_compat/mlx/linalg.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import mlx.core as mx +import mlx.core.linalg as mx_linalg + +from .._internal import clone_module, get_xp +from ..common import _linalg +from ._aliases import matmul, matrix_transpose, tensordot, vecdot # noqa: F401 + +__all__ = clone_module("mlx.core.linalg", globals()) + +EighResult = _linalg.EighResult +EigResult = _linalg.EigResult +QRResult = _linalg.QRResult +SlogdetResult = _linalg.SlogdetResult +SVDResult = _linalg.SVDResult + +# Wrap spec-compliant versions of linalg functions via common helpers +# eigh only works on CPU in MLX; pass stream=mx.cpu explicitly +def eigh(x: mx.array, /) -> _linalg.EighResult: + return EighResult(*mx_linalg.eigh(x, stream=mx.cpu)) +qr = get_xp(mx)(_linalg.qr) +slogdet = get_xp(mx)(_linalg.slogdet) +svd = get_xp(mx)(_linalg.svd) +cholesky = get_xp(mx)(_linalg.cholesky) +matrix_rank = get_xp(mx)(_linalg.matrix_rank) +pinv = get_xp(mx)(_linalg.pinv) +matrix_norm = get_xp(mx)(_linalg.matrix_norm) +vector_norm = get_xp(mx)(_linalg.vector_norm) +svdvals = get_xp(mx)(_linalg.svdvals) +diagonal = get_xp(mx)(_linalg.diagonal) +trace = get_xp(mx)(_linalg.trace) +cross = get_xp(mx)(_linalg.cross) +outer = get_xp(mx)(_linalg.outer) + +# MLX linalg.solve signature matches spec (no ambiguous stacked-vector +# behaviour like NumPy), so pass through directly. +solve = mx_linalg.solve + +# MLX linalg has no eig/eigvals returning complex — mark as unsupported. +# Raise a clear error rather than silently giving wrong results. +def eig(x, /): + raise NotImplementedError( + "eig() is not yet supported for MLX. " + "MLX linalg.eigh() is available for symmetric/Hermitian matrices." + ) + + +def eigvals(x, /): + raise NotImplementedError( + "eigvals() is not yet supported for MLX. " + "MLX linalg.eigh() is available for symmetric/Hermitian matrices." + ) + +_all = ["eig", "eigvals", "solve", "cross", "outer", + "matrix_transpose", "matmul", "tensordot", "vecdot"] + +__all__ = sorted( + set(__all__) + | set(_linalg.__all__) + | set(_all) +) + + +def __dir__() -> list[str]: + return __all__ diff --git a/tests/_helpers.py b/tests/_helpers.py index 17865aa0..83a25db7 100644 --- a/tests/_helpers.py +++ b/tests/_helpers.py @@ -2,13 +2,13 @@ import pytest -wrapped_libraries = ["numpy", "cupy", "torch", "dask.array"] +wrapped_libraries = ["numpy", "cupy", "torch", "dask.array", "mlx.core"] all_libraries = wrapped_libraries + [ "array_api_strict", "jax.numpy", "ndonnx", "sparse" ] def import_(library, wrapper=False): - pytest.importorskip(library) + pytest.importorskip(library.split(".")[0]) if wrapper: if 'jax' in library: # JAX v0.4.32 implements the array API directly in jax.numpy @@ -17,7 +17,10 @@ def import_(library, wrapper=False): if not hasattr(jax_numpy, "__array_api_version__"): library = 'jax.experimental.array_api' elif library in wrapped_libraries: - library = 'array_api_compat.' + library + if library == 'mlx.core': + library = 'array_api_compat.mlx' + else: + library = 'array_api_compat.' + library return import_module(library) diff --git a/tests/test_array_namespace.py b/tests/test_array_namespace.py index 311efc37..7c822ccc 100644 --- a/tests/test_array_namespace.py +++ b/tests/test_array_namespace.py @@ -42,7 +42,8 @@ def test_array_namespace(request, library, api_version, use_compat): elif library == "dask.array": assert namespace == array_api_compat.dask.array else: - assert namespace == getattr(array_api_compat, library) + compat_attr = "mlx" if library == "mlx.core" else library + assert namespace == getattr(array_api_compat, compat_attr) if library == "numpy": # check that the same namespace is returned for NumPy scalars diff --git a/tests/test_common.py b/tests/test_common.py index 85ed032e..50f8c23c 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -8,10 +8,10 @@ from array_api_compat import ( # noqa: F401 is_numpy_array, is_cupy_array, is_torch_array, is_dask_array, is_jax_array, is_pydata_sparse_array, - is_ndonnx_array, + is_ndonnx_array, is_mlx_array, is_numpy_namespace, is_cupy_namespace, is_torch_namespace, is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace, - is_array_api_strict_namespace, is_ndonnx_namespace, + is_array_api_strict_namespace, is_ndonnx_namespace, is_mlx_namespace, ) from array_api_compat import ( @@ -29,6 +29,7 @@ 'jax.numpy': 'is_jax_array', 'sparse': 'is_pydata_sparse_array', 'ndonnx': 'is_ndonnx_array', + 'mlx.core': 'is_mlx_array', } is_namespace_functions = { @@ -40,6 +41,7 @@ 'sparse': 'is_pydata_sparse_namespace', 'array_api_strict': 'is_array_api_strict_namespace', 'ndonnx': 'is_ndonnx_namespace', + 'mlx.core': 'is_mlx_namespace', } @@ -249,7 +251,7 @@ def test_asarray_cross_library(source_library, target_library, request): xfail(request, reason="The truth value of lazy Array Array(dtype=Boolean) is unknown") elif source_library == "ndonnx" and target_library == "numpy": xfail(request, reason="produces numpy array of ndonnx scalar arrays") - elif target_library == "ndonnx" and source_library in ("torch", "dask.array", "jax.numpy"): + elif target_library == "ndonnx" and source_library in ("torch", "dask.array", "jax.numpy", "mlx.core"): xfail(request, reason="unable to infer dtype") elif source_library == "jax.numpy" and target_library == "torch": @@ -305,6 +307,9 @@ def test_asarray_copy(library): a[0] = 0 assert b[0] == 0 + if library == "mlx.core": + return + # copy=None defaults to True when impossible a = asarray([1.0], dtype=xp.float32) assert a.dtype == xp.float32 diff --git a/tests/test_isdtype.py b/tests/test_isdtype.py index 92e45613..9f7ee6af 100644 --- a/tests/test_isdtype.py +++ b/tests/test_isdtype.py @@ -24,6 +24,21 @@ def _spec_dtypes(library): 'float32', 'float64', } + elif library == 'mlx.core': + # mlx does not support float64 or complex128 + return { + 'bool', + 'complex64', + 'float32', + 'int16', + 'int32', + 'int64', + 'int8', + 'uint16', + 'uint32', + 'uint64', + 'uint8', + } else: return { 'bool', diff --git a/tests/test_mlx.py b/tests/test_mlx.py new file mode 100644 index 00000000..271b657d --- /dev/null +++ b/tests/test_mlx.py @@ -0,0 +1,276 @@ +"""MLX-specific integration tests for array-api-compat.""" +import pytest + +try: + import mlx.core as mx +except ImportError: + pytestmark = pytest.skip(allow_module_level=True, reason="mlx not found") + +from array_api_compat import ( + array_namespace, + is_mlx_array, + is_mlx_namespace, + device, + to_device, +) +import array_api_compat.mlx as mlx_compat + + +# --- Detection helpers --- + +def test_is_mlx_array(): + x = mx.array([1.0, 2.0]) + assert is_mlx_array(x) + assert not is_mlx_array([1.0, 2.0]) + assert not is_mlx_array(1.0) + + +def test_is_mlx_namespace(): + assert is_mlx_namespace(mlx_compat) + assert is_mlx_namespace(mx) + import numpy as np + assert not is_mlx_namespace(np) + + +def test_array_namespace(): + x = mx.array([1.0, 2.0]) + xp = array_namespace(x) + assert is_mlx_namespace(xp) + + +# --- API version --- + +def test_array_api_version(): + assert mlx_compat.__array_api_version__ == "2025.12" + + +# --- Renamed functions --- + +def test_trig_renames(): + import math + x = mx.array([0.5]) + assert float(mlx_compat.acos(x)[0]) == pytest.approx(math.acos(0.5)) + assert float(mlx_compat.asin(x)[0]) == pytest.approx(math.asin(0.5)) + assert float(mlx_compat.atan(x)[0]) == pytest.approx(math.atan(0.5)) + + +def test_concat(): + a = mx.array([1, 2]) + b = mx.array([3, 4]) + result = mlx_compat.concat([a, b]) + assert list(result.tolist()) == [1, 2, 3, 4] + + +# --- Creation functions --- + +def test_arange(): + x = mlx_compat.arange(5) + assert list(x.tolist()) == [0, 1, 2, 3, 4] + + +def test_linspace(): + x = mlx_compat.linspace(0.0, 1.0, 5) + assert len(x.tolist()) == 5 + + +def test_zeros_ones_empty(): + assert mlx_compat.zeros((2, 3)).shape == (2, 3) + assert mlx_compat.ones((2, 3)).shape == (2, 3) + assert mlx_compat.full((2,), 7.0).shape == (2,) + + +def test_eye(): + e = mlx_compat.eye(3) + assert e.shape == (3, 3) + assert float(e[0, 0].tolist()) == 1.0 + assert float(e[0, 1].tolist()) == 0.0 + + +# --- asarray / astype --- + +def test_asarray(): + x = mlx_compat.asarray([1.0, 2.0], dtype=mx.float32) + assert is_mlx_array(x) + assert x.dtype == mx.float32 + + +def test_astype(): + x = mx.array([1.0, 2.0], dtype=mx.float32) + y = mlx_compat.astype(x, mx.float16) + assert y.dtype == mx.float16 + + +# --- iinfo shim --- + +def test_iinfo(): + info = mlx_compat.iinfo(mx.int32) + assert info.min == -(2**31) + assert info.max == 2**31 - 1 + assert info.bits == 32 + + info8 = mlx_compat.iinfo(mx.uint8) + assert info8.min == 0 + assert info8.max == 255 + + +def test_iinfo_from_array(): + x = mx.array([1], dtype=mx.int16) + info = mlx_compat.iinfo(x) + assert info.bits == 16 + + +# --- isdtype --- + +def test_isdtype(): + assert mlx_compat.isdtype(mx.float32, "real floating") + assert mlx_compat.isdtype(mx.int32, "signed integer") + assert mlx_compat.isdtype(mx.uint8, "unsigned integer") + assert mlx_compat.isdtype(mx.bool_, "bool") + assert not mlx_compat.isdtype(mx.float32, "integral") + + +# --- std / var (correction kwarg) --- + +def test_std_var(): + x = mx.array([1.0, 2.0, 3.0, 4.0]) + s = mlx_compat.std(x, correction=1.0) + v = mlx_compat.var(x, correction=1.0) + assert float(s.tolist()) == pytest.approx(float(mx.sqrt(mx.array(5.0 / 3.0)).tolist()), rel=1e-5) + _ = v # just confirm it runs + + +# --- sort / argsort with descending --- + +def test_sort_descending(): + x = mx.array([3, 1, 2]) + s = mlx_compat.sort(x, descending=True) + assert list(s.tolist()) == [3, 2, 1] + + +def test_argsort_descending(): + x = mx.array([3, 1, 2]) + idx = mlx_compat.argsort(x, descending=True) + assert list(idx.tolist()) == [0, 2, 1] + + +# --- reshape --- + +def test_reshape(): + x = mx.array([1, 2, 3, 4]) + y = mlx_compat.reshape(x, (2, 2)) + assert y.shape == (2, 2) + + +# --- permute_dims --- + +def test_permute_dims(): + x = mx.zeros((2, 3, 4)) + y = mlx_compat.permute_dims(x, (2, 0, 1)) + assert y.shape == (4, 2, 3) + + +# --- unstack --- + +def test_unstack(): + x = mx.array([[1, 2], [3, 4]]) + parts = mlx_compat.unstack(x, axis=0) + assert len(parts) == 2 + assert list(parts[0].tolist()) == [1, 2] + + +# --- cumulative_sum --- + +def test_cumulative_sum(): + x = mx.array([1, 2, 3]) + cs = mlx_compat.cumulative_sum(x, axis=0) + assert list(cs.tolist()) == [1, 3, 6] + + +def test_cumulative_sum_include_initial(): + x = mx.array([1, 2, 3]) + cs = mlx_compat.cumulative_sum(x, axis=0, include_initial=True) + assert list(cs.tolist()) == [0, 1, 3, 6] + + +# --- __array_namespace_info__ --- + +def test_namespace_info_capabilities(): + info = mlx_compat.__array_namespace_info__() + caps = info.capabilities() + assert caps["boolean indexing"] is True + assert "data-dependent shapes" in caps + + +def test_boolean_indexing_getitem(): + x = mx.array([1, 2, 3]) + mask = x < 3 + y = x[mask] + assert list(y.tolist()) == [1, 2] + +def test_namespace_info_dtypes(): + info = mlx_compat.__array_namespace_info__() + dtypes = info.dtypes() + assert "float32" in dtypes + assert "int32" in dtypes + assert "bool" in dtypes + # MLX has no float64 + assert "float64" not in dtypes + + +def test_namespace_info_default_dtypes(): + info = mlx_compat.__array_namespace_info__() + dd = info.default_dtypes() + assert dd["real floating"] == mx.float32 + + +def test_namespace_info_devices(): + info = mlx_compat.__array_namespace_info__() + devs = info.devices() + assert len(devs) >= 1 + + +# --- device / to_device --- + +def test_device(): + x = mx.array([1.0]) + d = device(x) + assert d is not None + + +def test_to_device(): + x = mx.array([1.0]) + d = device(x) + y = to_device(x, d) + assert is_mlx_array(y) + assert device(y) == d + + +# --- linalg submodule --- + +def test_linalg_available(): + assert hasattr(mlx_compat, "linalg") + + +def test_linalg_eigh(): + A = mx.array([[2.0, 1.0], [1.0, 2.0]]) + result = mlx_compat.linalg.eigh(A) + assert hasattr(result, "eigenvalues") + assert hasattr(result, "eigenvectors") + + +def test_linalg_eig_not_implemented(): + A = mx.array([[1.0, 0.0], [0.0, 1.0]]) + with pytest.raises(NotImplementedError): + mlx_compat.linalg.eig(A) + + +# --- fft submodule --- + +def test_fft_available(): + assert hasattr(mlx_compat, "fft") + + +def test_fft_basic(): + x = mx.array([1.0, 0.0, 0.0, 0.0]) + result = mlx_compat.fft.fft(x) + assert result.dtype == mx.complex64