Skip to content

Add array API standard aliases for existing ops#3678

Open
katlun-lgtm wants to merge 1 commit into
ml-explore:mainfrom
katlun-lgtm:array-api-aliases
Open

Add array API standard aliases for existing ops#3678
katlun-lgtm wants to merge 1 commit into
ml-explore:mainfrom
katlun-lgtm:array-api-aliases

Conversation

@katlun-lgtm

Copy link
Copy Markdown

Proposed changes

array.__array_namespace__() returns the mlx.core module, so array API conformance depends on mlx.core exposing 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:

array API name forwards to
acos, acosh, asin, asinh, atan, atanh arccos, arccosh, arcsin, arcsinh, arctan, arctanh
atan2 arctan2
pow power
bitwise_left_shift, bitwise_right_shift left_shift, right_shift

Each alias is a thin forwarder following the existing concat -> concatenate precedent, 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 on main; from_dlpack and other gaps remain for follow-ups.)

Checklist

  • I have read the CONTRIBUTING document
  • clang-format and black (the formatters configured in .pre-commit-config.yaml) report no changes on the modified files
  • Added a test (python/tests/test_ops.py::TestOps::test_array_api_aliases)
  • Built from source and ran python -m unittest test_ops locally — all 140 tests pass; the existing __array_namespace__ tests in test_array.py also pass

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant