jax.dlpack.from_dlpack

jax.dlpack.from_dlpack(external_array, device=None, copy=None)

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

如果未请求设备转移或复制，返回的 Array 与 external_array 共享内存。

参数:

• external_array – 具有 __dlpack__ 和 __dlpack_device__ 方法的数组对象。

• device (xla_client.Device | Sharding | None) – 可选的 设备，表示返回的数组应放置的设备。如果指定，则结果将被提交到该设备。如果未指定，结果数组将被解包到其来源的同一设备上。将 device 设置为与 external_array 来源不同的设备将需要复制，这意味着 copy 必须设置为 True 或 None。

• copy (bool | None) – 一个（可选的）布尔值，控制是否执行复制操作。如果 copy=True，则总是执行复制，即使解压到同一设备上。如果 copy=False，则从不执行复制，并且在必要时会引发错误。当 copy=None 时，如果需要进行设备转移，则可能会执行复制。

返回:

一个 jax.Array

备注

虽然 JAX 数组总是不可变的，但 dlpack 缓冲区不能被标记为不可变，并且外部进程有可能对其进行就地修改。如果一个 jax 数组是从 dlpack 缓冲区构造的，并且该缓冲区随后被就地修改，那么在操作相关的 JAX 数组时可能会导致未定义行为。