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.