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.

Takes ownership of the contents of x; leaves x in an invalid/deleted state.

Parameters

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