jax.dlpack.to_dlpack

目录

jax.dlpack.to_dlpack#

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

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

参数:

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

  • stream (int | Any | None) – 可选的依赖于平台的流，等待直到缓冲区准备好。这对应于 https://dmlc.github.io/dlpack/latest/python_spec.html 中记录的 __dlpack__ 的 stream 参数。

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

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

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

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

返回:

一个 DLPack PyCapsule 对象。

备注

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

目录