jax.dlpack module

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