- jax.dlpack.to_dlpack(x, take_ownership=False)#
Returns a DLPack tensor that encapsulates a
Takes ownership of the contents of
x; leaves x in an invalid/deleted state.
Any) – a
DeviceArray, on either CPU or GPU.
bool) – If
True, 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. If
False, JAX retains ownership of the buffer; it is undefined behavior if the DLPack consumer writes to a buffer that JAX owns.