jax.dlpack module

from_dlpack(external_array[, device, copy])	Returns a Array representation of a DLPack tensor.
is_supported_dtype(dtype)	Check if dtype is supported by jax.dlpack.