Struct burn_jit::tensor::JitTensor

pub struct JitTensor<R, E, const D: usize>
where
    R: Runtime,
    E: JitElement,
{
    pub client: ComputeClient<R::Server, R::Channel>,
    pub handle: Handle<R::Server>,
    pub shape: Shape<D>,
    pub device: R::Device,
    pub strides: [usize; D],
    /* private fields */
}

The basic tensor primitive struct.

Fields

client: ComputeClient<R::Server, R::Channel>

Compute client for the runtime.

handle: Handle<R::Server>

The buffer where the data are stored.

shape: Shape<D>

The shape of the tensor.

device: R::Device

The device of the tensor.

strides: [usize; D]

The strides of the tensor.

Implementations

impl<R, E, const D: usize> JitTensor<R, E, D>

where R: Runtime, E: JitElement,

pub fn new( client: ComputeClient<R::Server, R::Channel>, device: R::Device, shape: Shape<D>, handle: Handle<R::Server> ) -> Self

Create a new tensor with a contiguous memory layout.

pub fn to_client( &self, client: ComputeClient<R::Server, R::Channel>, device: R::Device ) -> Self

Change the context of the current tensor and return the newly transferred tensor.

pub fn copy(&self) -> Self

Copy the current tensor.

pub fn can_mut(&self) -> bool

Check if the tensor is safe to mutate.

pub fn assert_is_on_same_device(&self, other: &Self)

Assert that both tensors are on the same device.

pub fn is_contiguous(&self) -> bool

Check if the current tensor is contiguous.

Trait Implementations

impl<R, E, const D: usize> Clone for JitTensor<R, E, D>

where R: Runtime, E: JitElement,

fn clone(&self) -> Self

Returns a copy of the value.

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

Performs copy-assignment from source.

impl<R, E, const D: usize> Debug for JitTensor<R, E, D>

where R: Runtime, E: JitElement,

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.