Add compatibility layer for Apple mlx #162

Activity

1. lucascolley commented on Jul 24, 2024

   

   lucascolley

   on Jul 24, 2024

   Member

   More actions

   x-ref ml-explore/mlx#48
2. asmeurer commented on Jul 24, 2024

   

   asmeurer

   on Jul 24, 2024

   Member

   More actions

   @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
3. asmeurer commented on Aug 5, 2024

   

   asmeurer

   on Aug 5, 2024

   Member

   More actions

   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.
4. asmeurer commented on Aug 7, 2024

   

   asmeurer

   on Aug 7, 2024

   Member

   More actions

   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.

   🎉React with 🎉1Reacted by Ralf Gommers
5. asmeurer

   added

   Low Priority

   Medium Priority

   and removed

   Low Priority

   on Nov 15, 2024
6. asmeurer commented on Nov 15, 2024

   

   asmeurer

   on Nov 15, 2024

   Member

   More actions

   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.
7. aaishwarymishra commented on Jul 13, 2026

   

   aaishwarymishra

   on Jul 13, 2026

   Last edited by aaishwarymishra

   More actions

   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 :)
8. rgommers commented on Jul 14, 2026

   

   rgommers

   on Jul 14, 2026

   Member

   More actions

   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.
9. ev-br commented on Jul 17, 2026

   

   ev-br

   on Jul 17, 2026

   Member

   More actions

   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.
10. aaishwarymishra commented on Jul 18, 2026

    

    aaishwarymishra

    on Jul 18, 2026

    More actions

    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 :)