jax.dlpack.to_dlpack
jax.dlpack.to_dlpack#
- jax.dlpack.to_dlpack(x, take_ownership=False)[source]#
Returns a DLPack tensor that encapsulates a
DeviceArray
x.Takes ownership of the contents of
x
; leaves x in an invalid/deleted state.- Parameters
x (
Any
) – aDeviceArray
, on either CPU or GPU.take_ownership (
bool
) – IfTrue
, JAX hands ownership of the buffer to DLPack, and the consumer is free to mutate the buffer; the JAX buffer acts as if it were deleted. IfFalse
, JAX retains ownership of the buffer; it is undefined behavior if the DLPack consumer writes to a buffer that JAX owns.