Struct cubecl_core::frontend::Tensor

pub struct Tensor<T: CubeType> { /* private fields */ }

The tensor type is similar to the array type, however it comes with more metadata such as stride and shape.

Implementations

impl<T: CubeType> Tensor<T>

pub fn new() -> Self

Constructs a new Tensor.

impl<T: CubeType> Tensor<T>

pub fn stride<C: Index>(&self, _dim: C) -> UInt

Obtain the stride of input at dimension dim

pub fn shape<C: Index>(&self, _dim: C) -> UInt

Obtain the shape of input at dimension dim

pub fn len(&self) -> UInt

The length of the buffer representing the tensor.

Warning

The length will be affected by the vectorization factor. To obtain the number of elements, you should multiply the length by the vectorization factor.

pub fn rank(&self) -> UInt

Returns the rank of the tensor.

Trait Implementations

impl<T: CubeType> CubeType for Tensor<T>

type ExpandType = ExpandElementTyped<Tensor<T>>

fn init(context: &mut CubeContext, expand: Self::ExpandType) -> Self::ExpandType

Wrapper around the init method, necessary to type inference.

impl<C: CubeType> ExpandElementBaseInit for Tensor<C>

fn init_elem(_context: &mut CubeContext, elem: ExpandElement) -> ExpandElement

impl<E: CubeType, I: Into<UInt>> Index<I> for Tensor<E>

type Output = E

The returned type after indexing.

fn index(&self, _index: I) -> &Self::Output

Performs the indexing (container[index]) operation.

impl<E: CubeType, I: Into<UInt>> IndexMut<I> for Tensor<E>

fn index_mut(&mut self, _index: I) -> &mut Self::Output

Performs the mutable indexing (container[index]) operation.

impl<C: CubePrimitive> LaunchArg for Tensor<C>

type RuntimeArg<'a, R: Runtime> = TensorArg<'a, R>

The runtime argument for the kernel.

impl<C: CubePrimitive> LaunchArgExpand for Tensor<C>

fn expand( builder: &mut KernelBuilder, vectorization: Vectorization, ) -> ExpandElementTyped<Tensor<C>>

Register an input variable during compilation that fill the KernelBuilder.

fn expand_output( builder: &mut KernelBuilder, vectorization: Vectorization, ) -> ExpandElementTyped<Tensor<C>>

Register an output variable during compilation that fill the KernelBuilder.

impl<E: CubePrimitive> SliceOperator<E> for Tensor<E>

type Expand = ExpandElementTyped<Tensor<E>>

fn slice<Start: Index, End: Index>( &self, start: Start, end: End, ) -> &Slice<'_, E>

Return a read-only view of all elements comprise between the start and end index.

fn __expand_slice<Start: Index, End: Index>( context: &mut CubeContext, expand: Self::Expand, start: Start, end: End, ) -> ExpandElementTyped<Slice<'static, E>>

Expand function of SliceOperator::slice.

fn slice_mut<Start: Index, End: Index>( &mut self, start: Start, end: End, ) -> &mut SliceMut<'_, E>

Return a read-write view of all elements comprise between the start and end index.

fn __expand_slice_mut<Start: Index, End: Index>( context: &mut CubeContext, expand: Self::Expand, start: Start, end: End, ) -> ExpandElementTyped<SliceMut<'static, E>>

Expand function of SliceOperator::slice_mut.

fn slice_mut_unsafe<Start: Index, End: Index>( &self, start: Start, end: End, ) -> &mut SliceMut<'_, E>

Return a read-write view of all elements comprise between the start and end index.

fn __expand_slice_mut_unsafe<Start: Index, End: Index>( context: &mut CubeContext, expand: Self::Expand, start: Start, end: End, ) -> ExpandElementTyped<SliceMut<'static, E>>

Expand function of SliceOperator::slice_mut_unsafe.

fn as_slice(&self) -> &Slice<'_, E>

Reinterprete the current type as a read-only slice.

fn __expand_as_slice( context: &mut CubeContext, expand: Self::Expand, ) -> ExpandElementTyped<Slice<'static, E>>

Expand function of SliceOperator::as_slice.

fn as_slice_mut(&mut self) -> &mut SliceMut<'_, E>

Reinterprete the current type as a read-write slice.

fn __expand_as_slice_mut( context: &mut CubeContext, expand: Self::Expand, ) -> ExpandElementTyped<SliceMut<'static, E>>

Expand function of SliceOperator::as_slice_mut.

fn as_slice_mut_unsafe(&self) -> &mut SliceMut<'_, E>

Reinterprete the current type as a read-write slice.

fn __expand_as_slice_mut_unsafe( context: &mut CubeContext, expand: Self::Expand, ) -> ExpandElementTyped<SliceMut<'static, E>>

Expand function of SliceOperator::as_slice_mut_unsafe.

impl<T: Vectorized + CubeType> Vectorized for &Tensor<T>

fn vectorization_factor(&self) -> UInt

fn vectorize(self, _factor: UInt) -> Self

impl<T: Vectorized + CubeType> Vectorized for &mut Tensor<T>

fn vectorization_factor(&self) -> UInt

fn vectorize(self, _factor: UInt) -> Self

impl<T: Vectorized + CubeType> Vectorized for Tensor<T>

fn vectorization_factor(&self) -> UInt

fn vectorize(self, _factor: UInt) -> Self