Skip to content

Support indexing with any type which implmented __index__ - #3210

Merged
angeloskath merged 5 commits into
ml-explore:mainfrom
aisk:index
Mar 19, 2026
Merged

Support indexing with any type which implmented __index__#3210
angeloskath merged 5 commits into
ml-explore:mainfrom
aisk:index

Conversation

@aisk

@aisk aisk commented Mar 5, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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)

@angeloskath angeloskath mentioned this pull request Mar 16, 2026
Closed
4 tasks done
zcbenz
zcbenz reviewed Mar 17, 2026
View reviewed changes

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@aisk
aisk requested a review from zcbenz March 18, 2026 17:05
angeloskath
angeloskath requested changes Mar 19, 2026
View reviewed changes

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread python/src/indexing.cpp
}
if (!PyIndex_Check(obj.ptr())) {
return false;
}

@angeloskath angeloskath Mar 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@aisk

aisk commented Mar 19, 2026

Copy link
Copy Markdown
Contributor Author

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
angeloskath approved these changes Mar 19, 2026
View reviewed changes

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@angeloskath
angeloskath merged commit 5fa1a8d into ml-explore:main Mar 19, 2026
16 checks passed
@aisk
aisk deleted the index branch March 20, 2026 00:34
@zcbenz zcbenz mentioned this pull request Mar 20, 2026
Merged
@BrewTestBot BrewTestBot mentioned this pull request Apr 22, 2026
Closed
1 task done
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@angeloskath angeloskath angeloskath approved these changes
@zcbenz zcbenz Awaiting requested review from zcbenz

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Support of numpy scalar indexing

3 participants

@aisk @zcbenz @angeloskath