jax.dlpack.to_dlpack#

jax.dlpack.to_dlpack(x, stream=None, src_device=None, dl_device=None, max_version=None, copy=None)[源代码]#

返回一个 DLPack 张量，该张量封装了一个 Array x。

参数:

  • x (Array) – 一个 Array，在 CPU 或 GPU 上。

  • stream (int | Any | None | None) – 可选的平台相关流，等待直到缓冲区准备就绪。这对应于 __dlpack__ 文档中的 stream 参数，在 https://dmlc.github.io/dlpack/latest/python_spec.html 中进行了描述。

  • src_device (xla_client.Device | None | None) – CPU 或 GPU Device。

  • dl_device (tuple[DLDeviceType, int] | None | None) – DLPack 格式的 (dl_device_type, local_hardware_id) 元组，例如由 __dlpack_device__ 生成的。

  • max_version (tuple[int, int] | None | None) – 消费者（即 __dlpack__ 的调用者）支持的最大 DLPack 版本，以 (major, minor) 的 2 元组形式表示。此函数不保证返回 max_version 版本的胶囊。

  • copy (bool | None | None) – 一个布尔值，指示是否复制输入。如果 copy=True，则该函数必须始终复制。当 copy=False 时，该函数绝不能复制，并且在认为需要复制时必须引发错误。如果 copy=None，则该函数必须尽可能避免复制，但如果需要可以复制。

返回:

一个 DLPack PyCapsule 对象。

注意

虽然 JAX 数组始终是不可变的，但 DLPackManagedTensor 缓冲区不能标记为不可变，并且 JAX 外部的进程可以就地修改它们。如果从 JAX 数组派生的 DLPack 缓冲区被修改，则在使用关联的 JAX 数组时可能会导致未定义的行为。当 JAX 最终支持 DLManagedTensorVersioned (DLPack 1.0) 时，就可以指定缓冲区是只读的。