Is dlpack supported?

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?

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.

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.

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?

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

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.

How is this going? DId you need any additional dlpack support in mx.array?

For some reason this code fails for me (linux gpu):

import torch
import mlx.core as mx
import numpy as np

np.from_dlpack(mx.array([42]))

->

AttributeError                            Traceback (most recent call last)
[<ipython-input-18-26021316f7b0>](https://localhost:8080/#) in <cell line: 1>()
----> 1 np.from_dlpack(mx.array([42]))

AttributeError: type object 'mlx.core.array' has no attribute '__dlpack__'

mx.__version__=='0.2.0'

When running on a Mac machine, I get

In [4]: tensor = torch.from_dlpack(mx.array([42]))
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[4], line 1
----> 1 tensor = torch.from_dlpack(mx.array([42]))

File /opt/homebrew/lib/python3.10/site-packages/torch/utils/dlpack.py:101, in from_dlpack(ext_tensor)
     50 """from_dlpack(ext_tensor) -> Tensor
     51
     52 Converts a tensor from an external library into a ``torch.Tensor``.
   (...)
     98
     99 """
    100 if hasattr(ext_tensor, '__dlpack__'):
--> 101     device = ext_tensor.__dlpack_device__()
    102     # device is either CUDA or ROCm, we need to pass the current
    103     # stream
    104     if device[0] in (DLDeviceType.kDLGPU, DLDeviceType.kDLROCM):

AttributeError: 'array' object has no attribute '__dlpack_device__'

In [5]: np.from_dlpack(mx.array([42]))
Out[5]: array([42], dtype=int32)

showing numpy works but torch fails.

The linux distribution is an old version. You can try installing with Conda to get a new release for linux (we'll update the pypi one)
Ok looks like we need a __dlpack_device__

Just checking if it's implemented now?

Just saw it. Thanks!

Let us know if you have any trouble with that.

Navigation Menu

Repository navigation

Is dlpack supported? #1080

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Is dlpack supported? #1080

Description

Activity

awni commented on May 6, 2024

justinchuby commented on May 6, 2024

awni commented on May 6, 2024

justinchuby commented on May 6, 2024

justinchuby commented on May 6, 2024

awni commented on May 25, 2024

justinchuby commented on May 25, 2024

justinchuby commented on May 25, 2024

awni commented on May 25, 2024

justinchuby commented on Jun 17, 2024

justinchuby commented on Jun 17, 2024

awni commented on Jun 18, 2024

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions