Repository navigation
Add compatibility layer for Apple mlx #162
Description
Activity
x-ref ml-explore/mlx#48
@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
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.
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.
Reacted by Ralf GommersI 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.
Area Array API expectation MLX status today Current approach in PR Risk / open question Priority device()/to_device()device(to_device(x, dev)) == devNo real device ownership; unified memory Virtual id(array) -> deviceregistryIs 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 NumPyIs global monkeypatching too invasive? Can this be limited or upstreamed? Blocker float64on GPUSupported 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.
- 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.
- 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
- 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 :)
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.
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.
Hi I have made a tracking issue and added some of the spec details and results I got by running the tests at #450.
Can you maybe look at it, and maybe suggest next steps, as some of the future tests are dependent on them, should I run more tests and add more sections but some of the tests fail because of issues in MLX I discovered with my testing till now which maybe not that usefull I guess, or maybe we should try to fix these issues first before moving to other sections :)
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.