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 | None) – 可選的平台相關串流，用於等待直到緩衝區準備就緒。這對應於 https://dmlc.github.io/dlpack/latest/python_spec.html 中記錄的 __dlpack__ 的 stream 引數。

  • 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) 時，將可以指定緩衝區是唯讀的。