Struct PjRtBuffer

pub struct PjRtBuffer { /* private fields */ }

A buffer represents a view on a memory slice hosted on a device.

Implementations

impl PjRtBuffer

pub fn client(&self) -> &PjRtClient

The client that owns this buffer.

pub fn copy_to_device(&self, device: PjRtDevice<'_>) -> Result<PjRtBuffer>

Copy the buffer to a different device.

pub fn to_literal_sync(&self) -> Result<Literal>

Copy the buffer back to the host as a literal.

pub fn on_device_shape(&self) -> Result<Shape>

Retrieve the shape used by this buffer.

pub fn copy_raw_to_host_sync<T: ArrayElement>( &self, dst: &mut [T], offset: usize, ) -> Result<()>

Copy the data stored in a buffer to host memory in a blocking way.

Trait Implementations

impl Drop for PjRtBuffer

fn drop(&mut self)

Executes the destructor for this type.

impl FromRawBytes for PjRtBuffer

type Context = PjRtClient

fn from_raw_bytes( client: &Self::Context, ty: ElementType, dims: &[usize], bytes: &[u8], ) -> Result<Self>

fn read_npy<T: AsRef<Path>>(path: T, c: &Self::Context) -> Result<Self>

Reads a npy file and return the stored multi-dimensional array as a literal.

fn read_npz<T: AsRef<Path>>( path: T, c: &Self::Context, ) -> Result<Vec<(String, Self)>>

Reads a npz file and returns the stored multi-dimensional arrays together with their names.

fn read_npz_by_name<T: AsRef<Path>>( path: T, c: &Self::Context, names: &[&str], ) -> Result<Vec<Self>>

Reads a npz file and returns the stored multi-dimensional arrays for some specified names.