Repository navigation
MLX compatibility: Creation functions #452
Description
Activity
I wasn't able to run tests on this file mainly because,
- By default mlx arrays are on gpu stream and they dont support float64, many tests failed because of that.
- some tests were checking different datatypes the
complexbug also affected these tests. - No device parameter or method signature mismatch
- changed the title
[-]MLX compatibility: Creation methods[/-][+]MLX compatibility: Creation functions[/+]on Jul 19, 2026 No device parameter
This is systemic and is worth noting indeed. Does not block testing at this stage, not more than
cupywhich does not havedevicearguments either.By default mlx arrays are on gpu stream and they dont support float64, many tests failed because of that.
This is a nice observation and runs into a systemic problem in the test suite (cf data-apis/array-api-tests#440).
Meanwhile, can you run tests withset_default_deviceto CPU? https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.set_default_device.html
That'd be similar to JAX, data-apis/array-api-tests#368 (comment)
Your table is very useful, but is a tad hard to read because the "notes" column runs into the page width, could swap with "Test node" and "result" columns.
And it'd be useful to continue compiling a compatibility table for the rest of the spec in this same issue.
By default mlx arrays are on gpu stream and they dont support float64, many tests failed because of that.
In addition / as an alternative to setting the MLX default device to CPU, you may run the test suite with the environment variable
ARRAY_API_TESTS_SKIP_DTYPES=float64,complex128. Skipping float64 mostly works: if a test fails with this skip and MLX, it'd be helpful to double-check the test failure with the skip andarray_api_compat.numpy. If it still fails, the problem is with the test suite (please report it in data-apis/array-api-tests#431). If it passes with numpy and fails with MLX, then report the failure here.-
arange: MLX rejects the valid three-argument form when stop=None, e.g. arange(0, None, 1).
-
asarray: copy=False fails even when the input is already an MLX array with the requested dtype; MLX claims it cannot
avoid copying. -
empty_like: MLX does not accept the standard dtype keyword—even dtype=None. The error mentioning zeros_like also
suggests empty_like is internally aliased or misrouted. -
eye: two issues:
- Zero-sized dimensions such as eye(0) are rejected.
- Large valid diagonal offsets such as k=2147483648 overflow or fail MLX’s argument binding.
-
linspace: MLX does not implement the endpoint keyword.
-
meshgrid: MLX returns a list, but Array API 2025.12 requires a tuple.
-
ones_like: MLX does not accept the standard dtype keyword.
-
zeros_like: Same missing dtype keyword support.
what I found was the empty and empty_like was implemented on new release, but the docs wasn't updated so I marked them as missing, I will update that, out of 16, 8 passed and 8 failed.
For the
from_dlpackits's test is in a separate file and it didn't ran It tried to determine which device to use in two ways:- With copy=False, it reads x.device, but MLX arrays do not have that attribute.
- Otherwise, it call xp.array_namespace_info().devices(),but MLX does not implement array_namespace_info.
Even if I manually hardcode the MLX device test will fail as mlx does not have device parameter which it is passing.
The two passing tests were:
- test_dlpack.py::test_dlpack_device
- test_dlpack.py::test_dunder_dlpack
Only test_dlpack.py::test_from_dlpack failed.
Details
FAILED array_api_tests/test_creation_functions.py::test_arange - TypeError: arange(): incompatible function arguments. The following argument types are supported: FAILED array_api_tests/test_creation_functions.py::test_asarray_arrays - ValueError: Unable to avoid copy while creating an array as requested. FAILED array_api_tests/test_creation_functions.py::test_empty_like - TypeError: zeros_like(): incompatible function arguments. The following argument types are supp... FAILED array_api_tests/test_creation_functions.py::test_eye - ExceptionGroup: Hypothesis found 2 distinct failures. (2 sub-exceptions) FAILED array_api_tests/test_creation_functions.py::test_linspace - TypeError: linspace(): incompatible function arguments. The following argument types are suppor... FAILED array_api_tests/test_creation_functions.py::test_meshgrid - AssertionError: assert <class 'list'> == <class 'tuple'> FAILED array_api_tests/test_creation_functions.py::test_ones_like - TypeError: ones_like(): incompatible function arguments. The following argument types are suppo... FAILED array_api_tests/test_creation_functions.py::test_zeros_like - TypeError: zeros_like(): incompatible function arguments. The following argument types are supp...
-
Thanks @aaishwarymishra , this is very helpful! Quick triage:
{empty,zeros,ones}_likenot accepting adtypeargument: this looks like a simple enhancement in MLXeye(0): it's a simple edge case yes, an annoying one to run into, and a simple one to fix at the sourcearange(1, None, 1): likewise, this shouldn't be a hard fix in MLXlinspace(...., endpoint={True,False}: likewise
Of these,
linspaceI would very much prefer to not meddle with in the compat layer, as it's not intrinsically hard but is prone to small inconsistencies, and will need rather extensive testing to avoidmlxand_compat.mlxdivergencies.The other ones, one strategy could be to write small wrappers in a -compat fork/PR, make sure they pass the test suite, and then submit to MLX. When you do, feel free to link back here or ping me :-).
The wrappers ideally are temporary and are needed to see if there are other issues masked by a first problem.meshgridreturning lists not tuples: this was a deliberate change in 2025.12; NumPy changed in in 2.0, CuPy followed up in ENH: makebroadcast_arrays,meshgridreturn a tuple not list cupy/cupy#9582. I think there were no loud compaints about either change, which probably means the backwards compatibility effect is not serious in practice, so I'd advocate to just change it despite that it's a backcompat break.
These all are small-ish issues, would be nice to fix for user experience. In principle can fix in the compat layer, but even if we do, I'd very much prefer to have some strategy for adding them to MLX and dropping them from the compat layer.
-
asarray(copy=False)failure is pointing at some deeper problem. Would be nice to understand in more details, could you provide an standalone MRE? -
the lack of
.devicearray attribute is something which 1) we cannot work around in the -compat layer at all (we do not touch the array object); 2) is likely to be a tripping hazard going forward, as we improve the device support and work towards device-dependent dtype subsets, cf [RFC]: Updates for float16/bfloat16 and for dtypes that are lacking full support in libraries array-api#998
Working around the lack of
.deviceattribute is brittle, cfso it would be really great if MLX could add the attribute, even if it's not used internally.array-api-compat/src/array_api_compat/common/_helpers.py
Lines 781 to 792 in bbfe8c2
I'll link to ml-explore/mlx#3484 (comment) as the recent comment with the MLX devs's stance. I'd suggest starting with the upstream changes that are most concrete, uncontroversial, and won't require separate kernel's. The first four bullets in the comment above could be a good start?
eye(0)for example I'd expect to work well.meshgridis simple but don't start with a backwards compat break; it's trivial to work around in this repo..devicedefinitely leave to the end, I don't expect that to be well-received.@ev-br Hi, I have made some pr's mainly focusing on bugs and non breaking changes :)
Reacted by Evgeni BurovskiFor completeness, here are @aaishwarymishra 's upstream PRs:
Hi @ev-br I have a question I looked at the adding
dtypeargument in {empty,zeros,ones}_like one blocker that I noticed if I put it in standard position(a,dtype.device)is that many of sub methods that rely on these methods in mlx uses positional arguments and it could possibly be backward breaking too.
One solution could be to adddtypeafterdeviceorstreamin the case of the mlx like(a,device,dtype).I made an issue regarding it I am not sure if my wording was not that good cause I wasn't able to get a concrete answer on how to proceed
I would appreciate your opinion :)
Not sure I understand the concern. Comparing the signatures (am copy-pasting from your table above):
empty_like(a, /, *, stream=None) # the current MLX signature empty_like(x, /, *, dtype=None, device=None) # the Array API spec signatureI think making it
empty_like(x, /, *, dtype=None, stream=None)is backwards compatible because you're adding a keyword-only argument.In the upstream issue above you mention C++ signatures, but those are of course different since the whole positional/kwarg argument parsing works differently in C++. And as long as the python frontend correctly sends the dtype argument to C++, there should be no fundamental issue---what am I missing?
If the concern is how to structure updating the C++, what I'd do is I'd start with one function, say,
empty_like, add the new argument in its natural location, beforestream, update the call sites, send the PR and flag it to the reviewers whether they are OK with this new argument in the middle of the C++ signature, or whether they prefer it in the end.
Maybe also include the count of the places which will need to be updated for other functions, viagrep | wc -lor something like that.Additional statuses of other PRS and issues:
- dtype argument addition in _like creation methods Adding dtype parameter to creation methods like the zeroes_like and ones_like ml-explore/mlx#4001
Reacted by Evgeni BurovskiI don't think the copy inconsistency in
asarraycan be fixed upstream ml-explore/mlx#4036Thanks for the upstream updates @aaishwarymishra , that's really impressive!
I presume you build mlx from source when working on upstream PRs?
It would be great to rerun the comparison after your fixes (+the tip of array-api-tests which fixed the too-large-k issue in test_eye, but that's minor).
My impression is that the remaining failures should be mostly hard mlx limitations (device attribute, copy= of asarray etc) and divergences with backwards compatibility considerations (meshgrid).
Hi @ev-br I am waiting for the response on the ml-explore/mlx#4152 and ml-explore/mlx#4034 as after these 2 issues are resolved I think we would have covered most of the easy and solvable cases, before I re-run the tests again
Nice! Any chance you'd be interested in generalising what you did to, say, manipulation functions ?
Sure :)
Hi i re-ran tests
arange— rejects integers larger than signed 32-bit, such asstep=2**31, well I found the reason they use standard c++ int for int values for arange binding so nanobind tries to convert it to int32 and fails not sure if supporting int64 for parameters in the spec.asarray— raises forcopy=Falseknow issue.eye— GPU scatter does not supportuint64, so identity matrix creation fails for that dtype.linspace— does not support the Array APIendpointkeyword, there is an draft pr open but not sure how it will proceed Add endpoint parameter to linspace ml-explore/mlx#4184.meshgrid— returns alist; Array API 2025.12 requires atuplewell someone wanted to work on it so i have left it for now Returning a tuple instead of list for meshgrid to make it compliant with Array API ml-explore/mlx#4034.
Also in
hypothesis_helpers.py_sorted_dtypes = [d for category in _dtype_categories for d in category] _device_dtype_pairs = [ (dtype, device) for dtype in _sorted_dtypes for device in xp.__array_namespace_info__().devices() if dh.is_device_dlpack_compatible(device) if dh.is_dtype_device_compatible(dtype, device) ]
causes tests to fail as for now MLX does not support
xp.__array_namespace_info__so I have to remove it to run the tests probably missed the last time, I re-ran tests on fresh pull so i ran into again.causes tests to fail as for now MLX does not support xp.array_namespace_info
Also no device parameter. Another workaround would be to put it in a try block so it doesn't effect other parts of the test suit. See below some changes I made regarding this in my local branch:
https://github.com/prady0t/array-api-tests/pull/1/changes
Great progress!
So we're down to 5 failures out of 15 tests; two are being worked on, and one is not going to be fixed.The remaining ones,
arangeandeye, indeed look likeint32limitations: I triedn = xp.iinfo(xp.int32).max ; xp.arange(n-1, n+3)and it raises a TypeError, and so doesxp.eye(n + 1).shape. Would be best to report upstream.Could you please edit the top post @aaishwarymishra , with the current status and links to upstream issues/PRs? Including rejected ones, so that the top post links to ml-explore/mlx#4036 and so on.
Reacted by Aaishwarya Mishraok as of now only 2 tests are failing
FAILED array-api-tests/array_api_tests/test_creation_functions.py::test_asarray_arrays - ValueError: Unable to avoid copy while creating an array as requested. FAILED array-api-tests/array_api_tests/test_creation_functions.py::test_eye - ValueError: [scatter] GPU scatter does not yet support uint64 for the input or updates.Of which one is dependant on ml-explore/mlx#4328
I don't think for now the
from_dlpackbe made compliant as the test uses .device which mlx does not expose======================================== FAILURES ========================================= ____________________________________ test_from_dlpack _____________________________________ + Exception Group Traceback (most recent call last): | File "/Users/blanky/code/array-api/array-api-tests/array_api_tests/test_dlpack.py", line 68, in test_from_dlpack | dtype_device_pair=hh.device_dtype_pairs, | ^^^ | File "/Users/blanky/code/array-api/.venv/lib/python3.13/site-packages/hypothesis/core.py", line 2288, in wrapped_test | raise the_error_hypothesis_found | ExceptionGroup: Hypothesis found 3 distinct failures. (3 sub-exceptions) +-+---------------- 1 ---------------- | Traceback (most recent call last): | File "/Users/blanky/code/array-api/array-api-tests/array_api_tests/test_dlpack.py", line 80, in test_from_dlpack | tgt_device = x.device if copy is False else device | ^^^^^^^^ | AttributeError: 'mlx.core.array' object has no attribute 'device' | Failing test case: test_from_dlpack( | dtype_device_pair=(mlx.core.bool, None), | copy_kw={'copy': False}, | data=data(...), | ) | Draw 1: array([], dtype=bool) +---------------- 2 ---------------- | Traceback (most recent call last): | File "/Users/blanky/code/array-api/array-api-tests/array_api_tests/test_dlpack.py", line 91, in test_from_dlpack | y = xp.from_dlpack(x, **tgt_device_kw, **copy_kw) | TypeError: from_dlpack(): incompatible function arguments. The following argument types are supported: | 1. from_dlpack(x: DLPackCompatible, /, *, copy: bool | None = None) -> array | | Invoked with types: mlx.core.array, kwargs = { device: NoneType } | | ========== FAILING CODE SNIPPET: | y = from_dlpack(array([], dtype=bool), **tgt_device_kw, **copy_kw) with tgt_device_kw={'device': None} and copy_kw={} | ==================== | | Failing test case: test_from_dlpack( | dtype_device_pair=(mlx.core.bool, None), | copy_kw={}, | data=data(...), | ) | Draw 1: array([], dtype=bool) | Draw 2: True +---------------- 3 ---------------- | Traceback (most recent call last): | File "/Users/blanky/code/array-api/array-api-tests/array_api_tests/test_dlpack.py", line 94, in test_from_dlpack | assert y.device == x.device | ^^^^^^^^ | AttributeError: 'mlx.core.array' object has no attribute 'device' | | ========== FAILING CODE SNIPPET: | y = from_dlpack(array([], dtype=bool), **tgt_device_kw, **copy_kw) with tgt_device_kw={} and copy_kw={} | ==================== | | Failing test case: test_from_dlpack( | dtype_device_pair=(mlx.core.bool, None), | copy_kw={}, | data=data(...), | ) | Draw 1: array([], dtype=bool) | Draw 2: False +------------------------------------ ================================= short test summary info ================================= FAILED array-api-tests/array_api_tests/test_dlpack.py::test_from_dlpack - ExceptionGroup: Hypothesis found 3 distinct failures. (3 sub-exceptions) =============================== 1 failed, 2 passed in 0.43
So 3 tests in total are failing
Great progress indeed!
test_asarray_arrayswill just need to be skipped/xfailed in the test suite (via an xfails file similar to https://github.com/data-apis/array-api-compat/blob/main/torch-xfails.txt)test_eye: might be helpful to check whether building MLX with the linked MLX PR fixes the test? And comment on the upstream PR.
EDIT:
- for
test_dlpackcf Fix test failure if no device parameter array-api-tests#459. It's either fixable at the test suite level, or will keep being a known failure.
So, one known failure, two more being worked on.
test_asarray_arrayswill just need to be skipped/xfailed in the test suite (via an xfails file similar to https://github.com/data-apis/array-api-compat/blob/main/torch-xfails.txt)
I mean we can also add an wrapper for simply returning the original array when copy=False similar to how jax does even though they don't allow mutations similar to mlx.
test_eye: might be helpful to check whether building MLX with the linked MLX PR fixes the test? And comment on the upstream PR.
test_eyemight need to be skipped/xfailed too as its MLX limitation and I don't think they are open for workarounds ml-explore/mlx#4300 (comment)we can also add an wrapper for simply returning the original array when copy=False
We could yes. The question is whether we should. The less workarounds are there, the better, for many reasons. At this stage at least, I don't see a reason to rush a workaround.
test_eye might need to be skipped/xfailed too as its MLX limitation and I don't think they are open for workarounds ml-explore/mlx#4300 (comment)
Please please edit this in into the OP table, together with other updates from the last month.
Reacted by Aaishwarya Mishra
arange(start, /, stop=None, step=1, *, dtype=None, device=None)arange(start, stop, step, dtype=None, *, stream=None)stop=Nonenow accepted. Remaining failure: integer args larger than signed 32-bit (e.g.step=2**31) — the nanobind binding uses C++int, soxp.arange(n-1, n+3)withn = iinfo(int32).maxraisesTypeError. Still nodevice.test_arangeasarray(obj, /, *, dtype=None, device=None, copy=None)asarray(a, dtype=None, *, copy=None)copy=Falseraises even for an MLX array already of the requested dtype. Fix was proposed and closed unmerged. Nodevice;dtypenot keyword-only.test_asarray_arraysempty(shape, *, dtype=None, device=None)empty(shape, dtype=float32, *, stream=None)device.test_emptyempty_like(x, /, *, dtype=None, device=None)empty_like(a, /, *, dtype=None, stream=None)dtypenow accepted —empty_likealiaseszeros_like, so it inherited the fix. Still undocumented; nodevice.test_empty_likeeye(n_rows, n_cols=None, /, *, k=0, dtype=None, device=None)eye(n, m=None, k=0, dtype=float32, *, stream=None)eye(0)fixed upstream; the too-large-kfailure was a test-suite issue, fixed at the tip of array-api-tests. Remaining failure: GPU scatter does not supportuint64MLX limitation, so identity creation fails for that dtype.test_eyefrom_dlpack(x, /, *, device=None, copy=None)from_dlpack(x, /, *, copy=None)from_dlpack: MLX has nox.deviceand does not have device parameter in the signaturetest_dlpack.py::test_from_dlpackfull(shape, fill_value, *, dtype=None, device=None)full(shape, vals, dtype=None, *, stream=None)fill_valueis namedvals;dtypenot keyword-only; nodevice.test_fullfull_like(x, /, fill_value, *, dtype=None, device=None)full_like(a, vals, dtype=None, *, stream=None)fill_valueis namedvals;dtypenot keyword-only; nodevice.test_full_likelinspace(start, stop, /, num, *, dtype=None, device=None, endpoint=True)linspace(start, stop, num=50, dtype=float32, stream=None)endpointstill unsupported, soendpoint=Falsefails. Draft PR open. Nodevice.test_linspacemeshgrid(*arrays, indexing="xy")meshgrid(*arrays, sparse=False, indexing="xy", stream=None)list; Array API 2025.12 requires atuple.test_meshgridones(shape, *, dtype=None, device=None)ones(shape, dtype=float32, *, stream=None)dtypenot keyword-only; nodevice. Behavior passes.test_onesones_like(x, /, *, dtype=None, device=None)ones_like(a, /, *, dtype=None, stream=None)dtypeadded upstream as a keyword-only arg. Nodevice.test_ones_liketril(x, /, *, k=0)tril(x, k, *, stream=None)kaccepted positionally; spec makes it keyword-only. Behavior passes.test_triltriu(x, /, *, k=0)triu(x, k, *, stream=None)kaccepted positionally; spec makes it keyword-only. Behavior passes.test_triuzeros(shape, *, dtype=None, device=None)zeros(shape, dtype=float32, *, stream=None)dtypenot keyword-only; nodevice. Behavior passes.test_zeroszeros_like(x, /, *, dtype=None, device=None)zeros_like(a, /, *, dtype=None, stream=None)dtypeadded upstream as a keyword-only arg. Nodevice.test_zeros_likePlanned x-fails
test_asarray_arraystest_eyetest_from_dlpackdepends on fix: define dask reshape copy behavior #459