jax.dlpack module

jax.dlpack.from_dlpack(dlpack, backend=None)[source]

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)[source]

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.