jax.dlpack.to_dlpack#

jax.dlpack.to_dlpack(x, take_ownership=False, stream=None)[source]#

Returns a DLPack tensor that encapsulates a Array x.

Parameters:

  • x (Array) – a Array, on either CPU or GPU.

  • take_ownership (bool) – Deprecated for xla_extension_version greater or equal than 204. It is a no-op to set take_ownership. Will be deleted in 01/2024. For xla_extension_version less than 204, if True, the JAX buffer acts as if it were deleted.

  • stream (int | Any | None) – optional platform-dependent stream to wait on until the buffer is ready. This corresponds to the stream argument to __dlpack__ documented in https://dmlc.github.io/dlpack/latest/python_spec.html.