jax.dlpack module

from_dlpack(external_array[, device, copy])	Returns a Array representation of a DLPack tensor.
to_dlpack(x[, stream, src_device, ...])	Returns a DLPack tensor that encapsulates a Array x.