jax.dlpack.to_dlpack

Warning

This page was created from a pull request (#9655).

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) – 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.

previous

jax.dlpack.from_dlpack

next

jax.distributed module