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
impl PjRtClient
Sourcepub fn cpu() -> Result<Self>
pub fn cpu() -> Result<Self>
A CPU client, this can run computations on multiple CPUs at the same time.
Sourcepub fn gpu(memory_fraction: f64, preallocate: bool) -> Result<Self>
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
.
Sourcepub fn compile(&self, c: &XlaComputation) -> Result<PjRtLoadedExecutable>
pub fn compile(&self, c: &XlaComputation) -> Result<PjRtLoadedExecutable>
Compile a computation for this device, and return the executable.
Sourcepub fn device_count(&self) -> usize
pub fn device_count(&self) -> usize
The number of devices that this client has detected, e.g. the number of GPUs.
Sourcepub fn addressable_device_count(&self) -> usize
pub fn addressable_device_count(&self) -> usize
The number of devices that this client can use.
Sourcepub fn platform_name(&self) -> String
pub fn platform_name(&self) -> String
The name of the platform.
Sourcepub fn platform_version(&self) -> String
pub fn platform_version(&self) -> String
The version of the platform.
Sourcepub fn devices(&self) -> Vec<PjRtDevice<'_>>
pub fn devices(&self) -> Vec<PjRtDevice<'_>>
A list of devices attached to this client.
Sourcepub fn addressable_devices(&self) -> Vec<PjRtDevice<'_>>
pub fn addressable_devices(&self) -> Vec<PjRtDevice<'_>>
A list of devices that can be used by this client.
Sourcepub fn buffer_from_host_buffer<T: ArrayElement>(
&self,
data: &[T],
dims: &[usize],
device: Option<&PjRtDevice<'_>>,
) -> Result<PjRtBuffer>
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.
Sourcepub fn buffer_from_host_raw_bytes(
&self,
ty: ElementType,
data: &[u8],
dims: &[usize],
device: Option<&PjRtDevice<'_>>,
) -> Result<PjRtBuffer>
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.
Sourcepub fn buffer_from_host_literal(
&self,
device: Option<&PjRtDevice<'_>>,
literal: &Literal,
) -> Result<PjRtBuffer>
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
impl Clone for PjRtClient
Source§fn clone(&self) -> PjRtClient
fn clone(&self) -> PjRtClient
1.0.0 · Source§fn clone_from(&mut self, source: &Self)
fn clone_from(&mut self, source: &Self)
source
. Read more