dlpack confusion

[YuzhiLiu-ai]

import jax
import jax.dlpack as jdp
import jax.numpy as jnp
a = jnp.arange(10)
b = jdp.to_dlpack(a)
hasattr(b, '__dlpack__')

result : False

I want to transform from jax.array to dlpack and use mpi4py to send and receive it, because of the lack of attribute '_ _ dlpack _ ', it fails. The return of function 'jax.dlpack.to_dlpack' doesn't have the attribute, is it reasonable? if i want to transform dlpack to numpy array. the attribute ' _ dlpack _ _' is necessary for 'numpy.from_dlpack' function

[jakevdp]

Hi - this looks like an unfortunate design choice made by numpy: the numpy.from_dlpack function does not accept a dlpack object directly, but rather accepts an object with a __dlpack__() method that returns a dlpack buffer. You can work around this issue as follows:

import jax
import jax.dlpack as jdp
import jax.numpy as jnp
import numpy as np

class NumpyCompatibleWrapper:
  def __init__(self, dlpack_buffer):
    self.dlpack_buffer = dlpack_buffer

  def __dlpack__(self):
    return self.dlpack_buffer

a = jnp.arange(10)
b = jdp.to_dlpack(a)
b_wrapper = NumpyCompatibleWrapper(b)

c = np.from_dlpack(b_wrapper)
print(c)
# [0 1 2 3 4 5 6 7 8 9]

I don't think there's any reasonable way for us to fix this on the JAX side: we could, of course, add a __dlpack__ method to our dlpack buffers that would essentially return self, but that's a confusing design choice. Instead, I'd suggest you report this as an issue to numpy: it's quite confusing for a function called from_dlpack to not accept dlpack objects directly.

[jakevdp]

Alternatively, we could think about adding a __dlpack__ method to JAX arrays, so that you could do np.from_dlpack(a). That seems to be the API that the numpy developers have in mind here, rather than directly operating on dlpack buffers.

[jakevdp]

Tracking this here: #11496

[YuzhiLiu-ai]

Thank you !