Is dlpack supported?

[justinchuby]

We are looking to support mlx arrays in the ONNX in memory IR: https://onnxscript.ai/intermediate_representation/tensors.html

To create interoperability support we implement the __dlpack__ protocol (https://dmlc.github.io/dlpack/latest/) on the tensor adapters (https://onnxscript.ai/intermediate_representation/tensors.html#computation-with-different-frameworks).

I read https://ml-explore.github.io/mlx/build/html/usage/numpy.html and did a quick search but it doesn't seem mlx supports dlpack. Just wondering if there are any plans for it?

[awni]

MLX doesn't support the dlpack protocol yet. It does support the Python buffer protocol, could that work for you?

We could add dlpack support, I do not think it would be too difficult given that nanobind's nd::array supports it.

[justinchuby]

I think we can just wait for dlpack support, since it should be simpler to use. The tensor objects we create are simply wrappers around arrays from different frameworks, so it matters any target framework also needs to support the protocol.

[awni]

Do you need both support from dlpack and to dlpack? __dlpack__ gets you the capsule. But how do you go from dlpack to other framework? Or is that not needed?

[justinchuby]

I think allowing from_dlpack to be called on Arrays is enough. The implementation simply relays the call so the mlx array can be used by other frameworks: https://github.com/microsoft/onnxscript/blob/bc818e7ae4a4418dbc6a80da6d25311c7579e53d/onnxscript/ir/_core.py#L344

[justinchuby]

I imagine it’s the users responsibility to choose a target framework that can accept a dlpack compatible array. Of course it would be helpful for mlx to implement from_dlpack to accept arrays produced by other frameworks.