xla

Struct PjRtClient

Source
pub struct PjRtClient(/* private fields */);
Expand description

A client represents a device that can be used to run some computations. A computation graph is compiled in a way that is specific to a device before it can be run.

Implementations§

Source§

impl PjRtClient

Source

pub fn cpu() -> Result<Self>

A CPU client, this can run computations on multiple CPUs at the same time.

Source

pub fn gpu(memory_fraction: f64, preallocate: bool) -> Result<Self>

A GPU client, the memory requirements are limited by the specified memory_fraction and this memory can either be allocated dynamically or pre-allocated depending on preallocate.

Source

pub fn tpu(max_inflight_computations: usize) -> Result<Self>

A TPU client.

Source

pub fn compile(&self, c: &XlaComputation) -> Result<PjRtLoadedExecutable>

Compile a computation for this device, and return the executable.

Source

pub fn device_count(&self) -> usize

The number of devices that this client has detected, e.g. the number of GPUs.

Source

pub fn addressable_device_count(&self) -> usize

The number of devices that this client can use.

Source

pub fn platform_name(&self) -> String

The name of the platform.

Source

pub fn platform_version(&self) -> String

The version of the platform.

Source

pub fn devices(&self) -> Vec<PjRtDevice<'_>>

A list of devices attached to this client.

Source

pub fn addressable_devices(&self) -> Vec<PjRtDevice<'_>>

A list of devices that can be used by this client.

Source

pub fn buffer_from_host_buffer<T: ArrayElement>( &self, data: &[T], dims: &[usize], device: Option<&PjRtDevice<'_>>, ) -> Result<PjRtBuffer>

Transfer some data from the host to a PjRtBuffer stored on the target device. If the device is not specified, the default device is used. The source data is passed as a slice of the specified primitive type, as well as the dimensions. The dimensions have to match the number of elements in the source data, otherwise an error is returned.

Source

pub fn buffer_from_host_raw_bytes( &self, ty: ElementType, data: &[u8], dims: &[usize], device: Option<&PjRtDevice<'_>>, ) -> Result<PjRtBuffer>

Transfer some data from the host to a PjRtBuffer stored on the target device. If the device is not specified, the default device is used. The source data is passed as a slice of raw bytes, as well as the dimensions. The dimensions have to match the number of bytes in the source data, otherwise an error is returned.

Source

pub fn buffer_from_host_literal( &self, device: Option<&PjRtDevice<'_>>, literal: &Literal, ) -> Result<PjRtBuffer>

Transfer some data from the host to a PjRtBuffer stored on the target device. If the device is not specified, the default device is used. The source data is passed as a literal.

Trait Implementations§

Source§

impl Clone for PjRtClient

Source§

fn clone(&self) -> PjRtClient

Returns a copy of the value. Read more
1.0.0 · Source§

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source. Read more

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> CloneToUninit for T
where T: Clone,

Source§

unsafe fn clone_to_uninit(&self, dst: *mut u8)

🔬This is a nightly-only experimental API. (clone_to_uninit)
Performs copy-assignment from self to dst. Read more
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> Same for T

Source§

type Output = T

Should always be Self
Source§

impl<T> ToOwned for T
where T: Clone,

Source§

type Owned = T

The resulting type after obtaining ownership.
Source§

fn to_owned(&self) -> T

Creates owned data from borrowed data, usually by cloning. Read more
Source§

fn clone_into(&self, target: &mut T)

Uses borrowed data to replace owned data, usually by cloning. Read more
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = Infallible

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.