MLX compatibility: Manipulation functions

[aaishwarymishra]

Array API	MLX Analog	Status	Notes	Test Node	Result
broadcast_arrays(*arrays)	broadcast_arrays(*arrays, stream=None) -> list[array]	incompatible	Array API 2025.12 requires a tuple of arrays; MLX returns a list.	test_broadcast_arrays	Failed
broadcast_shapes(*shapes)	broadcast_shapes(*shapes)	compatible	MLX returns the expected broadcasted shape.	TestBroadcastShapes::test_broadcast_shapes	Passed
broadcast_to(x, /, shape)	broadcast_to(array, shape, *, stream=None)	compatible	Required broadcasting behavior passes. MLX adds the optional stream extension.	test_broadcast_to	Passed
concat(arrays, /, *, axis=0)	concat(arrays: list[array], axis: int | None = 0, *, stream=None)	compatible	Required behavior passes. MLX additionally accepts axis=None, positional axis, and stream.	test_concat	Passed
expand_dims(x, /, axis)	expand_dims(a, /, axis: int | Sequence[int], *, stream=None)	partial	Tuple-axis coverage passes, but MLX incorrectly rejects a valid negative insertion axis at the lower boundary.	TestExpandDims::test_expand_dims; TestExpandDims::test_expand_dims_tuples	Failed / Passed
flip(x, /, *, axis=None)	flip(a, /, axis: None | int | Sequence[int] = None, *, stream=None)	compatible	Passes for None, integer, and tuple axes. MLX additionally permits positional axis.	test_flip	Passed
moveaxis(x, source, destination, /)	moveaxis(a, /, source: int, destination: int, *, stream=None)	incompatible	Array API accepts integers or tuples for source and destination; MLX accepts only scalar integers and rejects valid calls such as moveaxis(x, (), ()).	test_moveaxis	Failed
permute_dims(x, /, axes)	permute_dims(a, /, axes: Sequence[int] | None = None, *, stream=None)	compatible	Required axis permutation behavior passes. MLX additionally permits omitted or None axes.	test_permute_dims	Passed
repeat(x, repeats, /, *, axis=None)	repeat(array: array, repeats: int, axis: int | None = None, *, stream=None)	partial	Array API allows scalar or 1D integer-array repeats, while MLX documents only scalar integers.	test_repeat	Failed
reshape(x, /, shape, *, copy=None)	reshape(a, /, shape: Sequence[int], *, stream=None)	partial	Tested behavior passes, but MLX lacks the standard copy keyword and therefore cannot accept all conforming calls.	test_reshape	Passed
roll(x, /, shift, *, axis=None)	roll(a, shift: int | tuple[int], axis: None | int | tuple[int] = None, /, *, stream=None)	partial	Tested behavior passes, but MLX makes shift and axis positional-only. Array API permits keyword shift and requires axis to be keyword-only.	test_roll	Passed
squeeze(x, /, axis)	squeeze(a, /, axis: None | int | Sequence[int] = None, *, stream=None)	compatible	Integer and tuple-axis behavior passes. MLX additionally allows omitted or None axes.	test_squeeze	Passed
stack(arrays, /, *, axis=0)	stack(arrays: list[array], axis: int | None = 0, *, stream=None)	compatible	Required behavior passes. MLX additionally accepts positional axis, axis=None, and stream.	test_stack	Passed
tile(x, repetitions, /)	tile(a, reps: int | Sequence[int], /, *, stream=None)	compatible	Required tiling behavior passes. The parameter-name difference is irrelevant because it is positional-only.	test_tile	Passed
unstack(x, /, *, axis=0)	unstack(x, /, *, axis: int = 0, stream=None) -> list[array]	incompatible	Array API requires a tuple of arrays; MLX returns a list, including [] instead of () for an empty axis.	test_unstack	Failed

[aaishwarymishra]

also found a bug in repeats data-apis/array-api-tests#450

[ev-br]

Thanks, this is helpful!

Tested behavior passes, but MLX lacks the standard copy keyword and therefore cannot accept all conforming calls.

hmmm, how many hypothesis examples do you use? IOW, did you set pytest ... --max-examples=...? if not, could you retry with --max-examples=1000 or more (this will make the test run longer).

Failed

Could you add a reproducible code snippet that the test suite prints out? For each failed test, preferably.

---

Quick triage:

• tuples vs lists (unsqueeze, broadcast_arrays`): this is similar to Returning a tuple instead of list for meshgrid to make it compliant with Array API ml-explore/mlx#4034 and Returning a tuple instead of list for meshgrid to make it compliant with Array API ml-explore/mlx#4034 (comment) seems to indicate that MLX would a accept a PR with a fix
• expand_dims with a negative axis: was it not fixed in the dev version? I.e. is there something beyond Fix segfault in expand_dims for out of bounds negative axes ml-explore/mlx#4021 worth reporting to MLX then.
• moveaxis : this is not about empty tuples, mlx does not support tuple axes at all: this is a moderate-size MLX enhancement then, hopefully they accept a PR

mx.moveaxis(a, (0, 1), (1, 0))
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[16], line 1
----> 1 mx.moveaxis(a, (0, 1), (1, 0))

TypeError: moveaxis(): incompatible function arguments. The following argument types are supported:
    1. moveaxis(a: array, /, source: int, destination: int, *, stream: Union[None, Stream, Device] = None) -> array

Invoked with types: mlx.core.array, tuple, tuple

• reshape copy= argument : I suspect this will be a "no", similar to asarray's argument, Fix copy bug in mlx.core.asarray ml-explore/mlx#4036 ; maybe they'll be open to accepting an argument and raising an error for unsupported values, similar to asarray.
• roll : maybe try to see if they accept a change to make signature Array API compatible; if not, this we can handle in the -compat layer.

[prady0t]

This table is quite useful! Thanks @aaishwarymishra! How did you generate it? (Or you did it manually!) Could you possibly generate a similar table for this info? I manually tried matching it via the docs, but it is incomplete, and I would love to have a cleaner table as above.

[aaishwarymishra]

@prady0t I think the keeping the scope small is the first thing, for example in this case I only ran the manipulation tests, got the methods that were failing, manually copy pasting the signature is quite time consuming, what I usually do let an AI, fetch the signatures from the array api docs and make an initial table for the mlx counter parts, and then I manually check there signatures from the array api docs and mlx docs to be sure they are correct and populate the results. Sometimes there are things that could get missed like in the case of creation methods, empty_like method wasn't listed in the MLX docs but its test case passed meaning it was undocumented I would have missed it if I had to run the whole test suit.
Also i would recommend you build the MLX from GitHub repo as some of the upstream pr's have fixed some of the issue that might have caused other tests to fail like recent addition of __complex__ now tests work fine with complex datatypes , float64 still fails but that's a different matter to be honest.

@ev-br for the case of --max-examples= no I haven't run it I was't aware that flag existed I will re-run tests and give update if anything changed.

[prady0t]

I think the keeping the scope small is the first thing

Agreed. I know you have been working hard on this, so I will keep that table as a reference for now. Let's focus on these ones first.

Also i would recommend you build the MLX from GitHub repo

Indeed. For me, building with the latest main brings down the error count close to 245 (this also includes my local fixes to the tests module). Does this number resonate with yours?

[ev-br]

float64 still fails but that's a different matter

Use $ ARRAY_API_TESTS_SKIP_DTYPES=float64 pytest ....

[aaishwarymishra]

I re-ran the tests with --max-examples=1000 no new test failed, results were similar.

Details

Array API Tests Module: mlx.core (0.32.1.dev20260812+3abd0fd6). API Version: 2025.12. Enabled Extensions: fft, linalg
rootdir: /Users/blanky/code/array-api/array-api-tests
configfile: pytest.ini
plugins: hypothesis-6.161.5, json-report-1.5.0, metadata-3.1.1
collected 13 items

array-api-tests/array_api_tests/test_manipulation_functions.py .F.F...F....F                                                   [100%]

============================================================== FAILURES ==============================================================
__________________________________________________ TestExpandDims.test_expand_dims ___________________________________________________

self = <array_api_tests.test_manipulation_functions.TestExpandDims object at 0x10a8f5590>

    @pytest.mark.unvectorized
>   @given(

        x=hh.arrays(dtype=hh.all_dtypes, shape=shared_shapes()),
        axis=shared_shapes().flatmap(
            # Generate both valid and invalid axis
            lambda s: st.integers(2 * (-len(s) - 1), 2 * len(s))
        ),
    )

array-api-tests/array_api_tests/test_manipulation_functions.py:127:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <array_api_tests.test_manipulation_functions.TestExpandDims object at 0x10a8f5590>, x = array(False, dtype=bool), axis = -2

    @pytest.mark.unvectorized
    @given(
        x=hh.arrays(dtype=hh.all_dtypes, shape=shared_shapes()),
        axis=shared_shapes().flatmap(
            # Generate both valid and invalid axis
            lambda s: st.integers(2 * (-len(s) - 1), 2 * len(s))
        ),
    )
    def test_expand_dims(self, x, axis):
        if axis < -x.ndim - 1 or axis > x.ndim:
            with pytest.raises(IndexError):
>               xp.expand_dims(x, axis=axis)
E               ValueError: [expand_dims] Axis -2 is out of bounds for array with 1 dimensions.
E               Failing test case: test_expand_dims(
E                   self=<array_api_tests.test_manipulation_functions.TestExpandDims object at 0x10a8f5590>,
E                   x=array(False, dtype=bool),
E                   axis=-2,
E               )

array-api-tests/array_api_tests/test_manipulation_functions.py:137: ValueError
___________________________________________________________ test_moveaxis ____________________________________________________________

    @pytest.mark.min_version("2023.12")
>   @given(x=hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_dims=1)), data=st.data())
                   ^^^

array-api-tests/array_api_tests/test_manipulation_functions.py:197:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

x = array([], dtype=bool), data = data(...)

    @pytest.mark.min_version("2023.12")
    @given(x=hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_dims=1)), data=st.data())
    def test_moveaxis(x, data):
        source = data.draw(
            st.integers(-x.ndim, x.ndim - 1) | xps.valid_tuple_axes(x.ndim), label="source"
        )
        if isinstance(source, int):
            destination = data.draw(st.integers(-x.ndim, x.ndim - 1), label="destination")
        else:
            assert isinstance(source, tuple)  # sanity check
            destination = data.draw(
                st.lists(
                    st.integers(-x.ndim, x.ndim - 1),
                    min_size=len(source),
                    max_size=len(source),
                    unique_by=lambda n: n if n >= 0 else x.ndim + n,
                ).map(tuple),
                label="destination"
            )

        repro_snippet = ph.format_snippet(f"xp.moveaxis({x!r}, {source!r}, {destination!r})")
        try:
>           out = xp.moveaxis(x, source, destination)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E           TypeError: moveaxis(): incompatible function arguments. The following argument types are supported:
E               1. moveaxis(a: array, /, source: int, destination: int, *, stream: StreamOrDevice = None) -> array
E
E           Invoked with types: mlx.core.array, tuple, tuple
E
E           ========== FAILING CODE SNIPPET:
E           xp.moveaxis(array([], dtype=bool), (), ())
E           ====================
E
E           Failing test case: test_moveaxis(
E               x=array([], dtype=bool),
E               data=data(...),
E           )
E           Draw 1 (source): ()
E           Draw 2 (destination): ()

array-api-tests/array_api_tests/test_manipulation_functions.py:218: TypeError
____________________________________________________________ test_repeat _____________________________________________________________
  + Exception Group Traceback (most recent call last):
  |   File "/Users/blanky/code/array-api/array-api-tests/array_api_tests/test_manipulation_functions.py", line 354, in test_repeat
  |     @given(
  |
  |   File "/Users/blanky/code/array-api/.venv/lib/python3.13/site-packages/hypothesis/core.py", line 2297, in wrapped_test
  |     raise the_error_hypothesis_found
  | ExceptionGroup: Hypothesis found 2 distinct failures. (2 sub-exceptions)
  +-+---------------- 1 ----------------
    | Traceback (most recent call last):
    |   File "/Users/blanky/code/array-api/array-api-tests/array_api_tests/test_manipulation_functions.py", line 386, in test_repeat
    |     out = xp.repeat(x, repeats, **kw)
    | TypeError: repeat(): incompatible function arguments. The following argument types are supported:
    |     1. repeat(array: array, repeats: int, axis: int | None = None, *, stream: StreamOrDevice = None) -> array
    |
    | Invoked with types: mlx.core.array, mlx.core.array
    |
    | ========== FAILING CODE SNIPPET:
    | xp.repeat(array([], dtype=bool),array([], dtype=int8), **kw) with kw={}
    | ====================
    |
    | Failing test case: test_repeat(
    |     x=array([], dtype=bool),
    |     kw={},
    |     data=data(...),
    | )
    | Draw 1 (repeats): array([], dtype=int8)
    +---------------- 2 ----------------
    | Traceback (most recent call last):
    |   File "/Users/blanky/code/array-api/array-api-tests/array_api_tests/test_manipulation_functions.py", line 415, in test_repeat
    |     ph.assert_array_elements("repeat", out=out_slice[start:end],
    |                                            ~~~~~~~~~^^^^^^^^^^^
    | ValueError: Slice indices must be integers or None.
    |
    | ========== FAILING CODE SNIPPET:
    | xp.repeat(array([False], dtype=bool),1, **kw) with kw={}
    | ====================
    |
    | Failing test case: test_repeat(
    |     x=array([False], dtype=bool),
    |     kw={},
    |     data=data(...),
    | )
    | Draw 1 (repeats): 1
    +------------------------------------
____________________________________________________________ test_unstack ____________________________________________________________

    @pytest.mark.min_version("2023.12")
>   @given(x=hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_dims=1)), data=st.data())
                   ^^^

array-api-tests/array_api_tests/test_manipulation_functions.py:587:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

x = array([], dtype=bool), data = data(...)

    @pytest.mark.min_version("2023.12")
    @given(x=hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_dims=1)), data=st.data())
    def test_unstack(x, data):
        axis = data.draw(st.integers(min_value=-x.ndim, max_value=x.ndim - 1), label="axis")
        kw = data.draw(hh.specified_kwargs(("axis", axis, 0)), label="kw")

        repro_snippet = ph.format_snippet(f"xp.unstack({x!r}, **kw) with {kw=}")
        try:
            out = xp.unstack(x, **kw)

>           assert isinstance(out, tuple)
E           assert False
E            +  where False = isinstance([], tuple)
E
E           ========== FAILING CODE SNIPPET:
E           xp.unstack(array([], dtype=bool), **kw) with kw={}
E           ====================
E
E           Failing test case: test_unstack(
E               x=array([], dtype=bool),
E               data=data(...),
E           )
E           Draw 1 (axis): 0
E           Draw 2 (kw): {}

array-api-tests/array_api_tests/test_manipulation_functions.py:596: AssertionError
====================================================== short test summary info =======================================================
FAILED array-api-tests/array_api_tests/test_manipulation_functions.py::TestExpandDims::test_expand_dims - ValueError: [expand_dims] Axis -2 is out of bounds for array with 1 dimensions.
FAILED array-api-tests/array_api_tests/test_manipulation_functions.py::test_moveaxis - TypeError: moveaxis(): incompatible function arguments. The following argument types are supported:
FAILED array-api-tests/array_api_tests/test_manipulation_functions.py::test_repeat - ExceptionGroup: Hypothesis found 2 distinct failures. (2 sub-exceptions)
FAILED array-api-tests/array_api_tests/test_manipulation_functions.py::test_unstack - assert False
==================================================== 4 failed, 9 passed in 31.29s ====================================================
(array-api) blanky@Blankys-MacBook-Air array-api % uv run pytest \
  array-api-tests/array_api_tests/test_data_type_functions.py::test_broadcast_arrays \
  array-api-tests/array_api_tests/test_data_type_functions.py::TestBroadcastShapes::test_broadcast_shapes \
  array-api-tests/array_api_tests/test_data_type_functions.py::test_broadcast_to --max-examples=1000
======================================================== test session starts =========================================================
platform darwin -- Python 3.13.13, pytest-9.1.1, pluggy-1.6.0
Environment variables:
----------------------
ARRAY_API_TESTS_MODULE = mlx.core
ARRAY_API_TESTS_SKIP_DTYPES = float64

Array API Tests Module: mlx.core (0.32.1.dev20260812+3abd0fd6). API Version: 2025.12. Enabled Extensions: fft, linalg
rootdir: /Users/blanky/code/array-api/array-api-tests
configfile: pytest.ini
plugins: hypothesis-6.161.5, json-report-1.5.0, metadata-3.1.1
collected 3 items

array-api-tests/array_api_tests/test_data_type_functions.py F..                                                                [100%]

============================================================== FAILURES ==============================================================
_______________________________________________________ test_broadcast_arrays ________________________________________________________

>   ???

/Users/blanky/code/array-api-tests/array_api_tests/test_data_type_functions.py:87:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

shapes = ((),), data = data(...)

>   ???
E   AssertionError: assert <class 'list'> == <class 'tuple'>
E    +  where <class 'list'> = type([array(False, dtype=bool)])
E
E   ========== FAILING CODE SNIPPET:
E   xp.broadcast_arrays(*arrays) with arrays = [array(False, dtype=bool)]
E   ====================
E
E   Failing test case: test_broadcast_arrays(
E       shapes=((),),
E       data=data(...),
E   )
E   Draw 1 (x1): array(False, dtype=bool)

/Users/blanky/code/array-api-tests/array_api_tests/test_data_type_functions.py:99: AssertionError
====================================================== short test summary info =======================================================
FAILED array-api-tests/array_api_tests/test_data_type_functions.py::test_broadcast_arrays - AssertionError: assert <class 'list'> == <class 'tuple'>
==================================================== 1 failed, 2 passed in 6.15s =====================================================
(

[aaishwarymishra]

@ev-br hi, after further investigation I found that TestExpandDims::test_expand_dims is not the rejecting the valid negative index error but the tests expects IndexError when an out of range index is given in axis but MLX throws the ValueError I am not sure it its test case problem or MLX

[rgommers]

The expand_dims spec says "IndexError – If provided an invalid axis, an IndexError should be raised." I think IndexError is the natural error to show a user, but it's not critical and the error type isn't "must" for that reason. It's also something that's probably uniformly done across a library and not specific to a single function. I'd suggest asking the MLX maintainers whether an out of bounds axis should raise IndexError everywhere and if a PR for that is welcome. If not, add an xfail/skip in the tests for this.