Ctrl+K

jax.dlpack.to_dlpack

Contents

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. It is a no-op to set take_ownership. Will be deleted in 01/2024.

  • 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.

Returns:

A dlpack PyCapsule object.

Note

While JAX arrays are always immutable, dlpack buffers cannot be marked as immutable, and it is possible for processes external to JAX to mutate them in-place. If a dlpack buffer derived from a JAX array is mutated, it may lead to undefined behavior when using the associated JAX array.

Contents