Adding MLX backend.

[Narsil]

What does this PR do?

Adds MLX backend.
MLX already has native safetensors support, but this was trivial to add.

Fixes # (issue) or description of the problem this PR solves.

[Narsil]

We can safely ignore the failure which is entirely unrelated to the PR (more recent clippy detects an issue with PyO3 itself, which will most likely be fixed, by simply upgrading PyO3)

[Vaibhavs10]

Dope! Thanks for putting this up! 🚀 Conceptually, this looks good to me.

That said, I'm not an expert in MLX so would appreciate if Awni could do a deeper review.

One qq: Can we potentially also test for persisting and loading quantised models as well? we don't need to be exhaustive here, but just testing for 4-bit should be okay IMO.

[Narsil]

Quantized is u8, therefore it should work out of the box (not sure how mlx persists quantized information in npx/npz, i'm guessing it's not saved)

[pcuenca]

Awesome! 🔥

[pcuenca]

Does this handle bfloat16?

[pcuenca]

Oh, I see in the tests below that it does :)

[Narsil]

Yup, using the same flax trick which is just a special named dtype.

[awni]

LGTM! Thanks for adding that!

I think using your built-in numpy.save is probably a good call from a maintenance standpoint.

But just a thought that this will involve a copy (on the load side as well). Currently we do a copy to get in and out of Numpy.

[Narsil]

But just a thought that this will involve a copy (on the load side as well). Currently we do a copy to get in and out of Numpy.

Good to know. The only thing really needed is frombuffer equivalent in order to get loading partial tensors working.
On the rust side, I use the slice information given by the user to create. a CPU buffer of the correct data within the tensor, and use frombuffer on given framework in order to send the correct data where it belongs.

If there's a copy from numpy-> MLX, that means that currently there are 2 copies created.

File -> Slice (local buffer inside rust) -> numpy.from_buffer (should be zero copy) -> mlx (new copy).

Is that correct ?

[awni]

If there's a copy from numpy-> MLX, that means that currently there are 2 copies created.

Yes that's correct. We can't avoid the copy from NumPy because we have to use specifically allocated memory to make it available to both the CPU and GPU.