cubecl_std/tensor/view/
idx.rs

1use cubecl_core::prelude::*;
2
3use crate::tensor::{View, ViewExpand, layout::Coordinates};
4
5impl<E: CubePrimitive, C: Coordinates, IO: Clone> CubeIndex for View<E, C, IO> {
6    type Output = E;
7    type Idx = C;
8}
9
10impl<E: CubePrimitive, C: Coordinates, IO: Clone> CubeIndexExpand for ViewExpand<E, C, IO> {
11    type Output = <E as CubeType>::ExpandType;
12    type Idx = <C as CubeType>::ExpandType;
13
14    fn expand_index(self, scope: &mut Scope, index: C::ExpandType) -> Self::Output {
15        self.__expand_read_method(scope, index)
16    }
17
18    fn expand_index_unchecked(self, scope: &mut Scope, index: C::ExpandType) -> Self::Output {
19        self.__expand_read_unchecked_method(scope, index)
20    }
21}
22
23impl<E: CubePrimitive, C: Coordinates> CubeIndexMut for View<E, C, ReadWrite> {}
24impl<E: CubePrimitive, C: Coordinates> CubeIndexMutExpand for ViewExpand<E, C, ReadWrite> {
25    fn expand_index_mut(self, scope: &mut Scope, index: C::ExpandType, value: Self::Output) {
26        self.__expand_write_method(scope, index, value)
27    }
28}