MLX Array API compatibility: upstream status, compatibility

[aaishwarymishra]

MLX Array API compatibility overview

This issue tracks the compatibility of MLX with the Python Array API
standard and coordinates potential changes across:

• MLX
• array-api-compat
• array-api-tests
• the Array API specification

Scope

• Array API version:
• MLX version/commit:
• array-api-tests commit:
• array-api-compat branch/commit:
• Platform: macos

Compatibility tracking:

1. Constants, Datatypes and Array Attributes MLX compatibility: Constants, Datatypes and Array Attributes #451
2. Creation Methods MLX compatibility: Creation functions #452
3. Manipulation Methods MLX compatibility: Manipulation functions #462
4. Statistical functions MLX compatibility: Statistical functions #463

---

The desired endgame

Upstream has a CI job which runs array-api-tests with a list of known failures.

The test suite needs to run with ARRAY_API_TESTS_SKIP_DTYPES=float64 environment variable, and --disable-data-dependent-shapes. For example,

ARRAY_API_TESTS_SKIP_DTYPES=float64 ARRAY_API_TESTS_MODULE=mlx.core pytest array_api_tests/test_searching_functions.py -v --disable-data-dependent-shapes --max-examples=1000

---

Process

Use this issue an umrella issue, open individual issues for groups of tests from the test suite, edit the links into this top post into the list above. (If you cannot edit it, ping @ev-br)

Problems/failures with individual functions, as uncovered by the test suite: discuss in the sub-issues. Keep this one for general blockers.

---

Known blockers

• missing __array_namespace_info__ : would be best to submit a PR upstream
• missing .device attribute of the array object : nothing to do about it

[prady0t]

Adding more info to this tracker:

• array-api-compat: Atleast __array_namespace_info__() is required as per standard which is not available in MLX. This has to be the first step addition to the compat library.
• array-api-tests: Changes needed (which are also needed for us to be able to at least run the test suite with MLX) are mentioned in this comment: None dtype comparison not supported in MLX array-api-tests#447 (comment), and the PR: None dtype comparison not supported in MLX array-api-tests#447 itself needs to be addressed.

With that said, running the test suite right now (skipping array_namespace_info and adding rest of the changes) looks something like this :

484 failed, 953 passed, 7 skipped, 64 warnings 

As suggested by @ev-br, trying to find a breadth-first summary of the failures (and this is from AI):

Root error / signature	Approx. affected tests	Fix count
TypeError: float() argument must be a string or a real number, not 'complex'	100+	One major root cause
Missing functions (AssertionError: ... is not defined / module has no attribute)	~60	Mostly adding wrappers/implementations
Signature mismatches (Argument 'axis' missing, wrong parameter names, etc.)	~70	One API compliance sweep
TypeError: incompatible function arguments	~15	Similar wrapper/signature fixes
Missing array methods/attributes (__index__, __complex__, device, mT, to_device)	~10	Small API additions
Missing info namespace (__array_namespace_info__)	~6	One implementation
Missing metadata (iinfo.bits, etc.)	~10	One implementation
Hypothesis found 2 distinct failures	~8	Needs investigation individually
Numerical assertion failures	~10	Independent algorithm bugs
RuntimeError loading Metal kernels	~40	One backend/kernel issue

Most of the test failures are because MLX has no __complex__ method on its array type, so every attempt to pull a Python complex scalar out of a complex-dtype MLX array falls back to trying __float__ instead (that's standard CPython behavior for the complex() constructor, not a bug in the test) and blows up.

Some of these failures, I believe, are also a fix in the array-api-tests itself. I tried changing complex(x) -> complex(float(xp.real(x)), float(xp.imag(x))) in some places and now we are left with 415 failures (all of the complex failures were not removed. It might make sense for MLX to xfail these but not sure yet!)

I also found a seg fault in this process (discovered via hypothesis):

import mlx.core as mlx

x = mlx.zeros(())
mlx.expand_dims(x, axis=(-3, -2))

Gives:

[1]    87881 segmentation fault 

Reporting this to mlx. The test suite is already uncovering bugs in mlx!

I'll continue to keep adding more info to this report.

Edit: Just saw ml-explore/mlx#3984, might help solve the complex issue mentioned above.

[aaishwarymishra]

Hi, I think it would be better to instead of running all the tests together make sub-issues for different categories for easier tracking :)

[prady0t]

Hi, I think it would be better to instead of running all the tests together make sub-issues for different categories for easier tracking

Hi, sorry, I don't understand. The above table has some categories. I haven't really worked on any failures yet, just pointed some out. The fix that I suggested was just to unblock the test suite and the ones that were bugs in MLX itself.

All the complex failures are gone now after the upstream PR merge in MLX! That means all the

float() argument must be a string or a real number, not 'complex'

errors are gone. Thanks for the effort @aaishwarymishra! Which means now we have

289 failed, 1146 passed

Hypothesis found another bug in MLX.

import mlx.core as mx
a = mx.zeros((0, 0))
r = mx.linalg.cholesky(a)
mx.eval(r)   # process aborts here

Output:

** On entry to  POTRF, parameter number  4 had an illegal value
libc++abi: terminating due to uncaught exception of type std::runtime_error: [Cholesky::eval_cpu] Cholesky decomposition failed with error code -4
[1]    31154 abort 

Reporting it too.

@ev-br Should we add __array_namespace_info__() in compat? Currently, I run the tests by:

if not hasattr(xp, "__array_namespace_info__"):
    pytest.skip(
        "__array_namespace_info__ is not defined in the array module",
        allow_module_level=True,
    )

[ev-br]

Thanks for pointing out failure categories!
My suggestion is to work through them in passes: you did the first pass, with categories, great. Next pass is to list individual failures per group of functions, in a style similar to #452, preferably with reproducible examples that the test suite emits. Then we can quickly triage them. And then for those which are easy-fix in MLX, send PRs to MLX. Observing @aaishwarymishra 's recent MLX PRs, it looks like they are receptive to PRs and are ready to merge them quickly. Opening issues upstream not sure, especially for small items like empty(0) or cholesky((0,0)).

[ev-br]

Should we add __array_namespace_info__() in compat?

If it helps working through unrelated items, yes why not. We won't merge MLX compat PRs not yet, so you could work off a -compat branch with __array_namespace_info__() if that makes it easier.

[prady0t]

list individual failures per group of functions

Makes sense. Thanks for the suggestion. Here's a tracker for all failing functions along with their failures. I'll add more to this table.

Function(s)	Failure	Minimal reproducer
__getitem__	Boolean indexing unsupported
__setitem__	Indexing semantics incompatible
__index__	Missing method	mx.asarray(1).__index__()
__pos__	Missing unary positive	+mx.asarray(1)
to_device	Missing method
device	Missing attribute	mx.asarray(1).device
mT	Missing attribute	mx.ones((2,2)).mT
__array_namespace__	Missing api_version parameter
__array_namespace_info__	Missing implementation	mx.__array_namespace_info__()
__dlpack__	Missing stream parameter
arange	Signature mismatch
asarray	Copy/dtype semantics
empty	Signature mismatch
empty_like	Missing dtype
eye	Missing k parameter
full	Signature mismatch
full_like	Missing fill_value
linspace	Missing num
meshgrid	Returns list instead of tuple	type(mx.meshgrid(mx.arange(3)))
ones	Signature mismatch
ones_like	Missing dtype
zeros	Signature mismatch
zeros_like	Missing dtype
astype	Signature mismatch
broadcast_arrays	Returns list instead of tuple
broadcast_shapes	Empty input handling	mx.broadcast_shapes()
can_cast	Incorrect dtype behavior
finfo	Constructor/signature issue
iinfo	Missing bits attribute	mx.iinfo(mx.int32).bits
result_type	Scalar handling differs
from_dlpack	Signature mismatch
fft, ifft, fftn, ifftn, rfft, irfft, rfftn, irfftn	Missing signature parameters
hfft, ihfft	Missing implementation	mx.fft.hfft(...)
fftfreq, rfftfreq	Missing d parameter / behavior
fftshift, ifftshift	Missing axes parameter
take_along_axis	Shape/signature mismatch
expand_dims	Invalid axis handling / segfault	mx.expand_dims(mx.zeros(()), axis=(-3, -2))
moveaxis	Missing destination
repeat	Behavioral failures
reshape	Signature mismatch
unstack	Incorrect behavior
broadcast_to	Signature mismatch
concat	Missing axis
diff	Missing axis
flip	Missing axis
permute_dims	Signature mismatch
roll	Signature mismatch
squeeze	Signature mismatch
stack	Missing axis
cholesky	Missing upper; aborts on empty (0,0)	mx.linalg.cholesky(mx.zeros((0,0)))
cross	Missing axis
det	Complex inputs unsupported
diagonal	Missing implementation	mx.linalg.diagonal(...)
eig	Wrong return type
eigh	Behavioral failures
eigvalsh	Incorrect handling
inv	Complex inputs unsupported
matmul	Integer dtype behavior differs
matrix_norm	Missing implementation	mx.linalg.matrix_norm(mx.eye(3))
matrix_power	Missing implementation	mx.linalg.matrix_power(mx.eye(2), 2)
matrix_rank	Missing implementation	mx.linalg.matrix_rank(mx.eye(3))
matrix_transpose	Missing implementation / wrong shape
outer	Missing implementation	mx.linalg.outer(...)
pinv	Missing rtol / behavioral failures
qr	Missing mode / behavioral failures
slogdet	Behavioral failures
solve	Shape/indexing failures
svd	Missing full_matrices / behavioral failures
svdvals	Missing implementation	mx.linalg.svdvals(...)
tensordot	Missing implementation	mx.linalg.tensordot(...)
trace	Missing implementation	mx.linalg.trace(...)
vecdot	Missing implementation	mx.linalg.vecdot(...)
vector_norm	Missing implementation	mx.linalg.vector_norm(...)
argmax	Wrong output dtype
argmin	Wrong output dtype
count_nonzero	Missing axis
nonzero	Missing implementation	mx.nonzero(...)
searchsorted	Missing implementation	mx.searchsorted(...)
argsort	Missing axis; behavioral failures
sort	Missing axis; behavioral failures
isin	Missing implementation	mx.isin(...)
unique_all	Missing implementation
unique_counts	Missing implementation
unique_inverse	Missing implementation
unique_values	Missing implementation
all, any, max, mean, min, prod, std, sum, var, cumulative_sum, cumulative_prod	Missing axis parameter and/or behavioral failures
acos, acosh, asin, asinh, cos, sin, tanh	Incorrect numerical/special-case behavior
atan2	Signature mismatch
bitwise_left_shift, bitwise_right_shift, bitwise_xor	Incorrect results
clip	Signature mismatch
copysign	Missing implementation	mx.copysign(...)
expm1	Incorrect complex special-case behavior
floor_divide	Incorrect results
hypot	Missing implementation	mx.hypot(...)
nextafter	Missing implementation	mx.nextafter(...)
remainder	Incorrect signed-zero behavior
sign	Incorrect NaN handling
signbit	Missing implementation	mx.signbit(...)

[ev-br]

Thanks, this is very helpful.

Would be great to continue filling in missing MWEs.
Some of these items, esp creation functions, are being worked on in #452

Quick triage:

• leave out things with data-dependent shapes (nonzero, unique_all etc)
• leave out device-related issues
• special cases for numerical functions (infs nans and such): we maintain lists of xfails for these for essentially all wrapped libraries, so can leave to the end
• several items look uncontroversial and free from backwards compatible considerations (__pos__, __index__, mT, hypot, empty inputs). Are you up to sending upstream PRs for these and related omissions? Would be best to coordinate with @aaishwarymishra
• simple signature tweaks, like renaming/permuting arguments can be done in the compat layer
• missing axis arguments: I actually wonder what it takes to implement it in MLX.
• some of the listed "signature mismatch" I'm not sure about. Take broadcast_to as an example. What is the mismatch? The MLX signature looks to be a proper superset of the Array API signature, modulo tuple vs sequence annotation (which can be ignored IMO):

broadcast_to(x: array,          /, shape: Tuple[int, ...]) → array
broadcast_to(a: scalar | array, /, shape: Sequence[int], *, stream: None | Stream | Device = None) → array

[ev-br]

Re missing axis argument in reductions: the docs list it, at least for some of them, eg. https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.any.html#mlx.core.any

Is this version-dependent, or are there some behaviour differences that the test suite catches?

[prady0t]

This has to be a Python version-specific issue. I ran the tests till now on Python 3.14 (inside a conda environment), and it failed for function signatures even though they were present in the nanobind C-extension function (nb::sig): https://github.com/ml-explore/mlx/blob/8056817bd1202477b6bed2b1f871da5960d64b28/python/src/ops.cpp#L1086

This kind of function signature is handled here:
https://github.com/data-apis/array-api-tests/blob/4b42cafee2e6a119e51a7178215aaba63e4663e9/array_api_tests/test_signatures.py#L265

And they are expected to raise a ValueError.

Here's the catch: on my local machines, Python 3.12.2 (from python.org), it passes. Some chance is, inspect.signature handling has been relaxed? I'm not sure, and I need to investigate it a bit more. Here's the CPython section that might be handling this: https://github.com/python/cpython/blob/115400b5090f14233b72b255e483b4485c868100/Lib/inspect.py#L2436

Should I raise an issue in CPython?

Also, with a workaround, I'm able to fix this by adding a condition in the test suite, and we are down to just 240 failures! All the axis not found and a few other function signature complaints are just inconsistencies.

The workaround (kuddos to Claude Code, I haven't tested it too well):

def _is_uninformative_signature(sig: Signature) -> bool:
    # Some Python versions (e.g. 3.14) don't raise ValueError for callables
    # backed by extension modules (e.g. nanobind, pybind11) that lack a
    # parseable text signature. Instead inspect.signature() falls back to a
    # generic "(*args, **kwargs)" signature which carries no real parameter
    # information, so we treat it the same as an uninspectable signature.
    params = list(sig.parameters.values())
    return len(params) == 2 and [p.kind for p in params] == [
        Parameter.VAR_POSITIONAL,
        Parameter.VAR_KEYWORD,
    ]

[ev-br]

Should I raise an issue in CPython?

Maybe. At the moment though, there seem to be too many moving parts: array-api-tests, mlx, nanobind and inspect.signature. It's really not yet clear where the problem is, not to me at least.
It'd be best to move the two first items out of the equation: cook up a dummy nanobind-wrapped module and see if the inspect problem reproduces in the same python version dependent way. If it does, there's still a question whether it's in inspect or in nanobind, but at least we know it's upstream not us.

[prady0t]

I don't think it's a nanobind issue. I created a minimal example and ran it on different versions of nanobind and Python in CI. See.

The results are consistent with nanobind v2.0.0 (2 years old), v2.13.0 (what mlx currently pins) and the latest.

It fails for these Python versions starting from 3.11.0 (The failure here means not raising the ValueError that our test suite expects):

3.11.9
3.11.13
3.12.3
3.12.10
3.13.0
3.13.9
3.14.0

Passes for:

3.11.0
3.11.8
3.12.0
3.12.2