jax.dlpack 模块

from_dlpack(external_array[, device, copy])	返回 DLPack 张量的 Array 表示。
is_supported_dtype(dtype)	检查 jax.dlpack 是否支持该数据类型。