Are there equivalents to numpy.tobytes(), numpy.frombuffer() in mlx?

[mzbac]

Currently, if I want to serialize the MLX tensor, I have to convert it to a numpy array, then use tobytes() to send it over. After that, I reconstruct the numpy array using numpy's frombuffer and finally convert it from numpy back to an MLX tensor. I am wondering if there are equivalent methods in MLX that can avoid the need for conversion to numpy, which could reduce latency when transferring tensors in a distributed environment.

[awni]

MLX implements the Python buffer protocol. So you can do:

>>> a =mx.array([1,2,3])
>>> x = memoryview(a)
>>> mx.array(x)
array([1, 2, 3], dtype=int32)

I'm curious though why not use mx.distributed? It will generally be faster as it keeps the communication ops inside the graph and avoids unnecessary copies.

[mzbac]

I am not very familiar with MPI and don't really know how to do it via mx.distribute for pipeline parallelism. Currently, I have only managed to use gRPC to send the tensor across the node. I will take a look at MPI and mx.distributed over the weekend to see if I can get it working.

[angeloskath]

Well currently for pipelining you would need send/recv which we haven't added yet, but they are fairly trivial to implement.

I will try to PR them before the weekend.

[mzbac]

@angeloskath @awni Thanks for helping with the question and clarification, really appreciate it ❤️