jax.dlpack.from_dlpack

内容

jax.dlpack.from_dlpack#

jax.dlpack.from_dlpack(external_array, device=None, copy=None)[source]#

返回 DLPack 张量的 Array 表示。

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

参数：:

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

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

  • copy (bool | None | None) – 可选的布尔值，控制是否执行复制。如果 copy=True，则始终执行复制，即使解包到同一设备上。如果 copy=False，则永远不会执行复制，如果需要，则会引发错误。当 copy=None 时，如果设备传输需要，则可能执行复制。

返回值:

一个 jax.Array

注意

虽然 JAX 数组始终是不可变的，但 dlpack 缓冲区无法被标记为不可变，并且外部进程有可能在 JAX 中就地修改它们。如果 jax Array 是从 dlpack 缓冲区构建的，并且该缓冲区后来在就地被修改，则在使用关联的 JAX 数组时可能会导致未定义的行为。

内容