Add array API standard aliases for existing ops#3678
Open
katlun-lgtm wants to merge 1 commit into
Open
Conversation
Adds spec-conformant aliases so mlx.core (returned by array.__array_namespace__()) exposes the array API names for ops that currently only exist under numpy-style names: acos, acosh, asin, asinh, atan, atanh -> arccos/arccosh/... atan2 -> arctan2 pow -> power bitwise_left_shift, bitwise_right_shift-> left_shift/right_shift Each alias forwards to the existing operation (same pattern as the existing concat -> concatenate alias) and is added to the ops docs. Adds test_array_api_aliases covering output equality with the originals and presence on the __array_namespace__() object. Partially addresses ml-explore#3484.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Proposed changes
array.__array_namespace__()returns themlx.coremodule, so array API conformance depends onmlx.coreexposing the array API standard names. Several elementwise ops are currently only available under their numpy-style names. This adds spec-conformant aliases that forward to the existing operations:acos,acosh,asin,asinh,atan,atanharccos,arccosh,arcsin,arcsinh,arctan,arctanhatan2arctan2powpowerbitwise_left_shift,bitwise_right_shiftleft_shift,right_shiftEach alias is a thin forwarder following the existing
concat->concatenateprecedent, is added to the ops documentation, and is covered by a new test (test_array_api_aliases) that checks output equality with the originals and presence on the__array_namespace__()object.This is a focused first step toward #3484. (
asarray/asarray(copy=)from that issue already landed onmain;from_dlpackand other gaps remain for follow-ups.)Checklist
clang-formatandblack(the formatters configured in.pre-commit-config.yaml) report no changes on the modified filespython/tests/test_ops.py::TestOps::test_array_api_aliases)python -m unittest test_opslocally — all 140 tests pass; the existing__array_namespace__tests intest_array.pyalso pass