array-api-compat icon indicating copy to clipboard operation
array-api-compat copied to clipboard

Add compatibility layer for Apple mlx

Open j-emberton opened this issue 1 year ago • 5 comments

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.

j-emberton avatar Jul 24 '24 18:07 j-emberton

x-ref ml-explore/mlx#48

lucascolley avatar Jul 24 '24 19:07 lucascolley

@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 avatar Jul 24 '24 19:07 asmeurer

Also there has been some upstream work in mlx to support the array API (e.g., https://github.com/ml-explore/mlx/pull/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 avatar Aug 05 '24 19:08 asmeurer

Based on the discussion at https://github.com/ml-explore/mlx/pull/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 avatar Aug 07 '24 22:08 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.

asmeurer avatar Nov 15 '24 00:11 asmeurer