MLX compatibility: indexing functions

[ev-br]

Towards #450.

test_take_along_axis fails because of a quirck of the MLX function signature: the axis argument is optional in the spec but MLX requires it:

This is array-api-strict:

>>> import array_api_strict as xp
TypeError: take_along_axis() takes 2 positional arguments but 3 were given
>>> xp.take_along_axis(xp.asarray([False], dtype=xp.bool), xp.asarray([], dtype=xp.int32), axis=None)
empty((0,), dtype=array_api_strict.bool)
>>> xp.take_along_axis(xp.asarray([False], dtype=xp.bool), xp.asarray([], dtype=xp.int32))
empty((0,), dtype=array_api_strict.bool)

and this is MLX:

>>> import mlx.core as mx
>>> mx.take_along_axis(mx.asarray([False], dtype=mx.bool_), mx.asarray([], dtype=mx.int32), None)
array([], dtype=bool)
>>> mx.take_along_axis(mx.asarray([False], dtype=mx.bool_), mx.asarray([], dtype=mx.int32))
Traceback (most recent call last):
  File "<python-input-30>", line 1, in <module>
    mx.take_along_axis(mx.asarray([False], dtype=mx.bool_), mx.asarray([], dtype=mx.int32))
    ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: take_along_axis(): incompatible function arguments. The following argument types are supported:
    1. take_along_axis(a: array, /, indices: array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array

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

[aaishwarymishra]

I think we need to add wrapper for indexing to use -1 axis default ml-explore/mlx#4357 or maybe explicitly use -1 axis value in the tests.

[ev-br]

No, why?
Let's see:

• the spec mandates axis=-1 default
• there was no default
• your PR added the axis=None default

So either

• mlx's axis=None value actually works as -1 , i.e. works along the last axis, or
• what you added is not spec-compliant.

If the former, then there's no need to do anything, it is spec-compliant after your PR.
If the latter, then it's best to quickly make it spec compliant in MLX, until a release freezes the default via backwards compatibility.

So which one is it?

[aaishwarymishra]

Ok so the axis=None in MLX behaves differently then the axis=-1, so for it to work we need to explicitly set axis=-1 .

I proposed them to add axis=-1, as default value but it would have been backward breaking but in this case as we had to previously supply the axis explicitly so all the old code written already had an explicit axis parameter anyways.

I proposed them to make axis=-1, but this is the response I got ml-explore/mlx#4357 (comment)

Maybe I wasn't able to explain them effectively :(

[ev-br]

Looking at https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.take_along_axis.html,
"axis=None means flatten" is a typical numpy look-alike, and it's neither array API compatible, nor is numpy's default.

Here's the numpy behavior (and a recent numpy is Array API compatible):

In [29]: a = np.arange(6).reshape(2, 3)

In [30]: a
Out[30]: 
array([[0, 1, 2],
       [3, 4, 5]])

In [31]: idx = np.asarray([[1, 0]])

In [32]: np.take_along_axis(a, idx)
Out[32]: 
array([[1, 0],
       [4, 3]])

In [33]: np.take_along_axis(a, idx, axis=-1)
Out[33]: 
array([[1, 0],
       [4, 3]])

In [34]: np.take_along_axis(a, idx.squeeze(), axis=None)   # note .squeeze, exception otherwise
Out[34]: array([1, 0])

In [35]: np.take_along_axis(a, idx, axis=0)
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
...
IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (1,2) (1,3)

Trying the same in MLX 0.32.1 (i.e., before your PR):

In [39]: ma = mx.arange(6).reshape(2, 3)

In [40]: midx = mx.asarray([[1, 0]])

In [41]: mx.take_along_axis(ma, midx, axis=-1)    # OK, agrees w/numpy
Out[41]: 
array([[1, 0],
       [4, 3]], dtype=int32)

In [42]: mx.take_along_axis(ma, midx)    # OK, was an error => no backwards compat
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
....
Invoked with types: mlx.core.array, mlx.core.array

In [43]: mx.take_along_axis(ma, midx, axis=None)    # OK agrees w/numpy
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[43], line 1
----> 1 mx.take_along_axis(ma, midx, axis=None)

ValueError: [take_along_axis] Indices of dimension 2 does not match array of dimension 1.

In [44]: mx.take_along_axis(ma, midx.squeeze(), axis=None)   # OK agrees w/numpy
Out[44]: array([1, 0], dtype=int32)

In [45]: mx.take_along_axis(ma, midx, axis=0)      # OK agrees w/numpy
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[45], line 1
----> 1 mx.take_along_axis(ma, midx, axis=0)

ValueError: [broadcast_shapes] Shapes (3) and (2) cannot be broadcast.

---

Therefore, I indeed think the best course of action is to send a follow-up PR changing to axis=-1 default, and explain that

• there is no default in any released version, hence the default value can change without any backwards compat impact
• the change of the default makes take_along_axis Array API compatible.

[aaishwarymishra]

sure I will make a follow up :)