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 数组。

返回类型:

数组

注意

虽然 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]