Add compatibility layer for Apple mlx

[j-emberton]

MLX is an array library for use on apple silicon. It is intended to closely follow the python Array API however some difference exist.

https://github.com/ml-explore/mlx?tab=readme-ov-file

It would be useful to add compatibility layers for this app to support users wishing to do array operations on apple gpu.

I am happy to lead this if I can be added as contributor.

[lucascolley]

x-ref ml-explore/mlx#48

[asmeurer]

@j-emberton this is definitely in scope for array-api-compat. If you want to work on this, the best thing to do would be to fork this repository and make a pull request.

Also note that partial support is OK to begin with, so if there's some function that's difficult to support, we can skip it for now. The development notes in the docs are a good place to start for contributing https://data-apis.org/array-api-compat/dev/index.html

[asmeurer]

Also there has been some upstream work in mlx to support the array API (e.g., ml-explore/mlx#1289), so you may want to synchronize with them to see what their plans are. If they are planning to add full support directly to mlx, that is preferable to adding support to the compat library. But if there are things they can't support easily because of backwards compatibility, then we may need a compat wrapper.

[asmeurer]

Based on the discussion at ml-explore/mlx#1289, the MLX folks are open to having full compatibility in MLX itself, so any work on upstream incompatibilities should go there. For now, let's just add helper functions for MLX to array-api-compat.

[asmeurer]

I don't know if the upstream is doing anything here, but it might be worth revisiting this if they aren't. MLX is a useful library.

[aaishwarymishra]

Area	Array API expectation	MLX status today	Current approach in PR	Risk / open question	Priority
device() / to_device()	device(to_device(x, dev)) == dev	No real device ownership; unified memory	Virtual id(array) -> device registry	Is virtual device tracking acceptable long-term? Any GC / identity edge cases?	Blocker
Boolean indexing (__getitem__)	Boolean masks must work	Native MLX raises on boolean indexing	Monkeypatch mx.array.__getitem__ and fallback to NumPy	Is global monkeypatching too invasive? Can this be limited or upstreamed?	Blocker
float64 on GPU	Supported where backend allows, or clearly handled	Not supported on MLX GPU	Document as limitation / future upstream work	Does this block compliance or just specific tests?	Blocker

@ev-br I have compiled 3 main blockers for the mlx implementation first.

1. As mlx is optimized for apple silicon with integrated cpu, gpu they don't support device I think we have to add a workaround for our side.
2. There is still no float64 support for the mlx gpu so we might have to skip some tests on gpu, Support for mlx.float64 ml-explore/mlx#799
3. There is limited support for the boolean indexing there was a issue for that [BUG] Boolean Indices not supported ml-explore/mlx#860, we may contact mlx if they want to try and implement it, or I can try to do it myself or make a wrapper.

I think before moving forward we should see how to handle these blockers, also for making a doc will notion be fine for future updates as these blockers were important and limited so i posted them here with the kind of workarounds I tried in my previous pr :)

[rgommers]

3. There is limited support for the boolean indexing there was a issue for that [BUG] Boolean Indices not supported ml-explore/mlx#860, we may contact mlx if they want to try and implement it, or I can try to do it myself or make a wrapper.

See array_api_extra.apply_where and how that is used in SciPy.

[ev-br]

in my previous pr

(for those getting notifications on this issue) The previous PR reference is #449.

think before moving forward we should see how to handle these blockers

I don't think we necessarily have to decide on these blockers first. On the contrary, I think it's better to start with simple(-er) things, and have a general overview of what is compliant and what isn't. Then have a clear idea of what passes the test suite, and proceed with triaging test suite failures. Essentially, this comment.