Re-introduce Array API Tests suite

[prady0t]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: Some parts were written using GitHub Copilot

Following #3526, we are in a much better shape to re-introduce the array-api-tests in CI.

[prady0t]

Please allow gh-action runs for this PR so we can see failures in CI.

[prady0t]

CC @ev-br

[ev-br]

One small suggestion below.

The first run is expected to have a few failures. Once we see them, we'll be able to produce a set of relevant skips.

[prady0t]

@ev-br We need to address this PR data-apis/array-api-tests#459

See error log in CI

Invoked with types: int, kwargs = { device: mlx.core.Device }
ERROR array_api_tests/test_dlpack.py - TypeError: zeros(): incompatible function arguments. The following argument types are supported:
    1. zeros(shape: int | Sequence[int], dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array

[ev-br]

Ah yes. Let's continue on the array-api-tests PR

[prady0t]

@ev-br The test suite now runs. Thanks!

I've tried adding a skip file in the latest commit (which may not be accurate, as I used AI to generate it); let's wait for another CI run.

[ev-br]

Thanks for the ping @prady0t!

IMO the tentative skips file shows that the compatibility level is very high.

To explain some context, which might not be immediately clear for those who do not deal with the Array API test suite regularly:

• the test suite itself mainly checks compatibility with the Array API spec. In particular, value testing it performs is rather limited. Checking correctness of implementations themselves is left to individual projects' test suites.
• the workflow itself is fairly typical for the Array API test suite. Here's the JAX analog: https://github.com/jax-ml/jax/blob/main/.github/workflows/jax-array-api.yml
• the test suite does property testing based on hypothesis. Therefore it is not entirely deterministic! The balance is in its --max-examples pytest switch: increasing the value improves stability and makes the run time longer.
• the skips/xfails file, skips-file= argument and xfail.txt added here is the test suite way of recording known failures. The test suite itself is cross-project, and runs unmodified for torch, jax, numpy, ndonnx and so on; therefore it does not use the standard @xfail / @skip decorators; instead, each project carries its own skipfile.

Having a non-empty skips file is actually expected: for example, here are skipfiles for jax, numpy and ndoonx.

Quickly looking through the skips this PR adds, several classes of failures might be somewhat lower priority:

• multiple array_api_tests/test_special_cases.py failures are what torch labels as "module:nans and infs"--- they probe limiting cases of mathematical functions, exp(+inf) = inf, exp(-inf)=0 etc; No known library implements them all fully;
• array_api_tests/test_linalg.py failures flag mostly missing kernels, especially in the complex plane. Presumably quite hard to implement;
• array_api_tests/test_operators_and_elementwise_functions.py flag several issues with elementwise kernels in the complex plane (potentially related to [BUG] mx.abs / mx.sqrt / mx.log on complex64 are silently wrong on GPU across a third of float32's range #4344 or similar);
• array_api_tests/test_has_names.py failures flag missing top-level functions (signbit, nextafter etc).

We could certainly refactor the xfails file to visually group failures.
We would also be happy to help triaging the failures if it's helpful.
Would you prefer us to open separate issues for individual failures, the groups of failures, just keep them grouped and annotated in the skips file, or something else entirely?

[prady0t]

Can we have another CI run for this PR?

[zcbenz]

Would you prefer us to open separate issues for individual failures, the groups of failures, just keep them grouped and annotated in the skips file, or something else entirely?

I think having a skips file is enough. When working on CUDA backend we also had a big skips file and we just gradually fixed the items in the list.

[ev-br]

Thank you @zcbenz !

@prady0t : It'd be best to group the entries in the skips file, in a style similar to e.g. https://github.com/data-apis/array-api-compat/blob/main/torch-xfails.txt :

• special cases go last (are least priority)
• extensions (linalg, fft)
• missing functions go together
• missing complex kernels
• missing arguments (dtype, descending etc)
• loss of precision in the complex plane

Let me know if you want to me to help with grouping and triage.

[prady0t]

It'd be best to group the entries in the skips file

Just to minimise some manual work, I again asked AI to regroup according to your suggestions. The last xfail file didn't work due to inline comments; those have now been removed in the latest commit. Feel free to change the group/order of the tests in the xfail file if you disagree.

[prady0t]

I've also added the xfail option (instead of the earlier skip option) with the strict option. This will help us catch when some tests start passing with future changes.

[prady0t]

@ev-br In the latest runs, apart from some missing/extra xfails, we have some crashes

See the run: https://github.com/ml-explore/mlx/actions/runs/34209419794/job/102015510711

These are the tests:

 Worker crashes
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1[1]
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1[None]
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[1]
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[None]
array_api_tests/test_creation_functions.py::test_eye

Would it make sense to put them in a skip file?

[ev-br]

A typical use case for both skips and xfails file is if something segfaults; then this something goes to the skips so that pytest does not even run it. Trying locally with mx.__version__ == 0.32.2

Worker crashes
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1[1]
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1[None]
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[1]
array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[None]

these all passed for me.

array_api_tests/test_creation_functions.py::test_eye

this fails with

E           ValueError: [scatter] GPU scatter does not yet support uint64 for the input or updates.
E           
E           ========== FAILING CODE SNIPPET:
E           xp.eye(1, None, **kw) with kw = {'dtype': mlx.core.uint64}
E           ====================

This is IIUC #4300 (comment) be best to add a link as a comment in the skips file, I'd think?

[ev-br]

It'd be great to

• ask the AI helper to group the failures by the test module, for each category
  (so that e.g. all "other failures" of test_manipulation_functions.py are visually together)
• manually triage "other failures" and "missing arguments" listings a bit, with one-liner descriptions
  and possibly repro examples from the test suite output. Where a failure has been triaged to be a wontfix,
  include a link to the discussion

As an example, I worked through failures in test_{creation,manipulation}_functions locally:

# copy=False is not supported, https://github.com/ml-explore/mlx/pull/4036
array_api_tests/test_creation_functions.py::test_asarray_arrays

# uint64 is not supported, https://github.com/ml-explore/mlx/issues/4300
array_api_tests/test_creation_functions.py::test_eye

# the spec mandates an IndexError for invalid axis, MLX raises a ValueError
array_api_tests/test_manipulation_functions.py::TestExpandDims::test_expand_dims

# MLX does not accept tuple axes: `mx.moveaxis(mx.eye(3), (), ())`
array_api_tests/test_manipulation_functions.py::test_moveaxis

# only integer `repeats` are accepted
array_api_tests/test_manipulation_functions.py::test_repeat

# copy=False is not accepted
array_api_tests/test_manipulation_functions.py::test_reshape

# returns a list instead of a tuple
array_api_tests/test_manipulation_functions.py::test_unstack

# scalar arguments not accepted
array_api_tests/test_operators_and_elementwise_functions.py::test_binary_with_scalars_real[atan2]

# min,max arguments are required
array_api_tests/test_operators_and_elementwise_functions.py::test_clip

While that's a bit of work, I think it'd be helpful to see which other failures are showstoppers (copy=False) and which are small discrepancies (atan2 scalar arguments) or even possibly just small bugs (clip requiring min/max).

EDIT: checked test_searching_functions.py, too. Two classes of failures: return dtype incompatible with the spec, and a missing argument

# out.dtype is unsigned if in.dtype is unsigned, instead of the default integer dtype
FAILED array_api_tests/test_searching_functions.py::test_argmax
FAILED array_api_tests/test_searching_functions.py::test_argmin

# missing `sorter` argument
FAILED array_api_tests/test_searching_functions.py::test_searchsorted
FAILED array_api_tests/test_searching_functions.py::test_searchsorted_with_scalars