Skip to content

Add initial MLX support with the creation wrappers#453

Open
aaishwarymishra wants to merge 2 commits into
data-apis:mainfrom
aaishwarymishra:mlx
Open

Add initial MLX support with the creation wrappers#453
aaishwarymishra wants to merge 2 commits into
data-apis:mainfrom
aaishwarymishra:mlx

Conversation

@aaishwarymishra

@aaishwarymishra aaishwarymishra commented Jul 27, 2026

Copy link
Copy Markdown

This pull request introduces initial support for the MLX backend in the array_api_compat package. The main changes add MLX-specific modules, implement compatibility wrappers for core array functions, and update the build and development configuration to include MLX. Below are the most important changes:

For #452

MLX Backend Implementation:

  • Added the new array_api_compat/mlx backend with __init__.py, _aliases.py, and _typing.py to support MLX arrays and provide compatibility wrappers for common array operations.
  • Implemented MLX-compatible versions of arange, asarray, empty_like, eye, meshgrid, ones_like, and zeros_like in _aliases.py, handling MLX-specific quirks and edge cases.
  • Defined MLX-specific type aliases (Array, DType, Device) in _typing.py for consistent type handling.
  • Updated __init__.py to expose all relevant symbols and ensure proper module structure for the MLX backend.

Build and Development Environment:

  • Updated meson.build and pyproject.toml to include MLX sources and conditionally require the mlx package on macOS (sys_platform == 'darwin'). [1] [2]

Copilot AI review requested due to automatic review settings July 27, 2026 00:06

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@aaishwarymishra

aaishwarymishra commented Jul 27, 2026

Copy link
Copy Markdown
Author

After adding temporary compatibility wrappers and rerunning test_creation_functions.py, I found a few additional MLX limitations:

  • arange dtype: MLX defaults integer inputs to float32, while the Array API expects int32. The wrapper uses int32 for integer-only inputs and float32 otherwise.
  • Large integer arange: mx.arange rejects some values above 2**31 - 1. Empty and single-element ranges can be handled, but nontrivial int64 and uint64 ranges still fail.
  • eye with int64/uint64: MLX GPU scatter does not support these dtypes, so the wrapper creates the matrix in int32 and casts it.
  • asarray(copy=False): Returning an existing compatible MLX array avoids a copy error, but the aliasing test fails because MLX does not support 0-D assignment with x[...] = value.
  • empty_like, ones_like, zeros_like: Wrappers use x.dtype when dtype=None.
  • meshgrid: The wrapper converts MLX’s list output to the required tuple.
  • linspace: Still fails because MLX does not support endpoint. No wrapper was added due to possible numerical differences.

With:

ARRAY_API_TESTS_SKIP_DTYPES=float64,complex64,int64,uint64

the suite reports 14 passed, 2 failed:
Note: Some times we got lucky with arange and other methods due to hypothesis gives different values which sometimes passes the tests

  1. test_asarray_arrays: unsupported 0-D Ellipsis assignment.
  2. test_linspace: missing endpoint support.

NOTE: MLX does not implement eq for comparing None with datatype, also we skipped complex completely due to problem with complex() but tests doesn't allow completely skipping complex, I think tests can be improved here

@ev-br

ev-br commented Jul 27, 2026

Copy link
Copy Markdown
Member

Thanks!

Going forward, it would be helpful to list reproducible examples from the test suite failures, which the test suite prints on a failure.

complex() but tests doesn't allow completely skipping complex, I think tests can be improved here

Please open an issue in the test suite, with an MRE. If it's a test suite problem, it should be reproducible with numpy or array-api-strict, I suppose.

MLX does not implement eq for comparing None with datatype

This is one of things to propose adding to MLX; even if they decline, we'll have an answer.
Similar for other fixes, large and small.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants