jax.dlpack module

jax.dlpack.from_dlpack(dlpack, backend=None)

Returns a DeviceArray representation of a DLPack tensor dlpack.

The returned DeviceArray shares memory with dlpack.

Parameters:

• dlpack – a DLPack tensor, on either CPU or GPU.
• backend – experimental, optional: the platform on which dlpack lives.

jax.dlpack.to_dlpack(x)

Returns a DLPack tensor that encapsulates a DeviceArray x.

The DLPack shares memory with x.

Parameters:x (DeviceArray) – a DeviceArray, on either CPU or GPU.