jax.numpy.from_dlpack

内容

jax.numpy.from_dlpack#

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

通过 DLPack 构造 JAX 数组。

JAX 实现 numpy.from_dlpack()。

参数：

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

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

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

返回:

输入缓冲区的 JAX 数组。

返回类型:

数组

注意

虽然 JAX 数组始终是不可变的，但 dlpack 缓冲区无法标记为不可变，并且 JAX 之外的进程可能能够就地修改它们。如果从 dlpack 缓冲区构造 JAX 数组而不进行复制，并且稍后就地修改源缓冲区，则在使用关联的 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]
内容