jax.dlpack module

from_dlpack(dlpack)	Returns a DeviceArray representation of a DLPack tensor.
to_dlpack(x[, take_ownership])	Returns a DLPack tensor that encapsulates a DeviceArray x.