cubecl_std/tensor/view/
idx.rs1use 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}