jax.dlpack 模块

jax.dlpack 模块#

from_dlpack(external_array[, device, copy])

返回一个 DLPack 张量的 Array 表示。

to_dlpack(x[, stream, src_device, ...])

返回一个封装了 Array x 的 DLPack 张量。