jax.numpy.from_dlpack#

jax.numpy.from_dlpack(x, /, *, device=None, copy=None)[源代码]#

通过 DLPack 构建 JAX 数组。

JAX 对 numpy.from_dlpack() 的实现。

参数:

  • x (Any) – 一个通过 __dlpack__ 和 __dlpack_device__ 方法实现 DLPack 协议的对象，或者 CPU 或 GPU 上的旧版 DLPack 张量。

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

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

返回:

输入缓冲区的 JAX 数组。

返回类型:

Array

注意

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

示例

通过 DLPack 在 NumPy 和 JAX 之间传递数据

>>> import numpy as np
>>> rng = np.random.default_rng(42)
>>> x_numpy = rng.random(4, dtype='float32')
>>> print(x_numpy)
[0.08925092 0.773956   0.6545715  0.43887842]
>>> hasattr(x_numpy, "__dlpack__")  # NumPy supports the DLPack interface
True

>>> import jax.numpy as jnp
>>> x_jax = jnp.from_dlpack(x_numpy)
>>> print(x_jax)
[0.08925092 0.773956   0.6545715  0.43887842]
>>> hasattr(x_jax, "__dlpack__")  # JAX supports the DLPack interface
True

>>> x_numpy_round_trip = np.from_dlpack(x_jax)
>>> print(x_numpy_round_trip)
[0.08925092 0.773956   0.6545715  0.43887842]