jax.dlpack.to_dlpack

jax.dlpack.to_dlpack(x, take_ownership=False)

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) – a DeviceArray, on either CPU or GPU.

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