Struct TensorArg

pub struct TensorArg<P: ReduceDType, RA: ReduceArgs, Tag> { /* private fields */ }

Implementations

impl<P: ReduceDType, RA: ReduceArgs> TensorArg<P, RA, Input>

pub fn new_input(_state: &RA::State<P>) -> Self

pub fn __expand_new_input( _scope: &mut Scope, state: <RA::State<P> as CubeType>::ExpandType, ) -> TensorArgExpand<P, RA, Input>

impl<P: ReduceDType, RA: ReduceArgs> TensorArg<P, RA, Output>

pub fn new_output(_state: &mut RA::State<P>) -> Self

pub fn __expand_new_output( _scope: &mut Scope, state: <RA::State<P> as CubeType>::ExpandType, ) -> TensorArgExpand<P, RA, Output>

Trait Implementations

impl<P: ReduceDType, RA: ReduceArgs, Tag> CubeType for TensorArg<P, RA, Tag>

type ExpandType = TensorArgExpand<P, RA, Tag>

fn into_mut(scope: &mut Scope, expand: Self::ExpandType) -> Self::ExpandType

Wrapper around the init method, necessary to type inference.

impl<P: ReduceDType, RA: ReduceArgs> VirtualTensorOperations<<P as ReduceDType>::In> for TensorArg<P, RA, Input>

fn read(&self, _index: u32) -> Line<E>

Read the tensor at the given index.

fn write(&self, _index: u32, _value: Line<E>)

Write the tensor at the given index.

fn shape(&self, _axis: u32) -> u32

Get the shape of the tensor at the given axis.

fn stride(&self, _axis: u32) -> u32

Get the stride of the tensor at the given axis.

fn rank(&self) -> u32

Get the rank of the tensor.

impl<P: ReduceDType, RA: ReduceArgs> VirtualTensorOperations<<P as ReduceDType>::Out> for TensorArg<P, RA, Output>

fn read(&self, _index: u32) -> Line<E>

Read the tensor at the given index.

fn write(&self, _index: u32, _value: Line<E>)

Write the tensor at the given index.

fn shape(&self, _axis: u32) -> u32

Get the shape of the tensor at the given axis.

fn stride(&self, _axis: u32) -> u32

Get the stride of the tensor at the given axis.

fn rank(&self) -> u32

Get the rank of the tensor.