Support indexing with any type which implmented __index__

[aisk]

Proposed changes

Any type which implemented __index__ can be used to indexing mlx array, so we can use numpy scalar to indexing array.

bool is not supported also it implemented __index__ (actually, bool is a subclass of int) because it's to quirky.

Close: #2710

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[zcbenz]

Can you fix the linting error? The code looks good to me.

[angeloskath]

Numpy arrays and mlx arrays do not implement __index__ so no need for the extra checks.

[angeloskath]

I think this should simply be return PyIndex_Check(obj.ptr()) && !nb::isinstance<nb::bool_>(obj);

[aisk]

Hi @angeloskath thank you for your review, but numpy has implmeneted __index__:

>>> np.int64(21).__index__()
21
>>> np.array([1,2]).__index__()
Traceback (most recent call last):
  File "<python-input-11>", line 1, in <module>
    np.array([1,2]).__index__()
    ~~~~~~~~~~~~~~~~~~~~~~~~~^^
TypeError: only integer scalar arrays can be converted to a scalar index

Even the np.array have __index__ slots on it, but it will raise error while calling, so we can't simplify the codes by just check PyIndex_Check(obj.ptr()) && !nb::isinstance<nb::bool_>(obj).

And there is also a test case for this:

https://github.com/ml-explore/mlx/blob/main/python/tests/test_array.py#L2065

Simplify the code will break it:

ERROR: test_setitem_with_boolean_mask (__main__.TestArray.test_setitem_with_boolean_mask)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/asaka/Codes/mlx/python/tests/test_array.py", line 2073, in test_setitem_with_boolean_mask
    mx.arange(1000).reshape(10, 10, 10)[mask_np] = 0
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^
TypeError: only integer scalar arrays can be converted to a scalar index

[angeloskath]

Hm sorry about the noise. It 's so weird cause I explicitly tested it but I guess it was too late and I confused the exceptions.