tenferro_tensor_core/rank.rs
1use crate::{Result, ShapeVec, StrideVec, ValidationError};
2use std::fmt::Debug;
3
4/// Rank contract for tensor metadata shapes and strides.
5///
6/// # Examples
7///
8/// ```rust
9/// use tenferro_tensor_core::{Rank, TensorRank};
10///
11/// let shape = <Rank<2> as TensorRank>::shape_from_vec(vec![2, 3].into())?;
12/// assert_eq!(shape.as_ref(), &[2, 3]);
13/// # Ok::<(), tenferro_tensor_core::ValidationError>(())
14/// ```
15pub trait TensorRank: private::Sealed + Clone + Copy + Debug + Eq + Send + Sync + 'static {
16 /// Static rank when known at compile time.
17 ///
18 /// # Examples
19 ///
20 /// ```rust
21 /// use tenferro_tensor_core::{DynRank, Rank, TensorRank};
22 ///
23 /// assert_eq!(DynRank::RANK, None);
24 /// assert_eq!(Rank::<2>::RANK, Some(2));
25 /// ```
26 const RANK: Option<usize>;
27
28 /// Shape representation for this rank.
29 ///
30 /// # Examples
31 ///
32 /// ```rust
33 /// use tenferro_tensor_core::{DynRank, TensorRank};
34 ///
35 /// let shape: <DynRank as TensorRank>::Shape = vec![2, 3].into();
36 /// assert_eq!(shape.as_ref(), &[2, 3]);
37 /// ```
38 type Shape: Clone + Debug + PartialEq + Eq + AsRef<[usize]>;
39
40 /// Stride representation for this rank.
41 ///
42 /// # Examples
43 ///
44 /// ```rust
45 /// use tenferro_tensor_core::{DynRank, TensorRank};
46 ///
47 /// let strides: <DynRank as TensorRank>::Strides = vec![1, 2].into();
48 /// assert_eq!(strides.as_ref(), &[1, 2]);
49 /// ```
50 type Strides: Clone + Debug + PartialEq + Eq + AsRef<[isize]>;
51
52 /// Convert a dynamic shape vector into this rank's shape representation.
53 ///
54 /// # Examples
55 ///
56 /// ```rust
57 /// use tenferro_tensor_core::{Rank, TensorRank};
58 ///
59 /// let shape = <Rank<1> as TensorRank>::shape_from_vec(vec![4].into())?;
60 /// assert_eq!(shape.as_ref(), &[4]);
61 /// # Ok::<(), tenferro_tensor_core::ValidationError>(())
62 /// ```
63 ///
64 /// # Errors
65 ///
66 /// Returns [`ValidationError::RankMismatch`] when the shape length does
67 /// not equal the compile-time rank.
68 fn shape_from_vec(shape: ShapeVec) -> Result<Self::Shape>;
69
70 /// Convert this rank's shape representation into a dynamic shape vector.
71 ///
72 /// # Examples
73 ///
74 /// ```rust
75 /// use tenferro_tensor_core::{Rank, TensorRank};
76 ///
77 /// let shape = <Rank<2> as TensorRank>::shape_from_vec(vec![2, 3].into())?;
78 /// assert_eq!(<Rank<2> as TensorRank>::shape_into_vec(shape).as_slice(), &[2, 3]);
79 /// # Ok::<(), tenferro_tensor_core::ValidationError>(())
80 /// ```
81 fn shape_into_vec(shape: Self::Shape) -> ShapeVec;
82
83 /// Convert a dynamic stride vector into this rank's stride representation.
84 ///
85 /// # Examples
86 ///
87 /// ```rust
88 /// use tenferro_tensor_core::{Rank, TensorRank};
89 ///
90 /// let strides = <Rank<2> as TensorRank>::strides_from_vec(vec![1, 2].into())?;
91 /// assert_eq!(strides.as_ref(), &[1, 2]);
92 /// # Ok::<(), tenferro_tensor_core::ValidationError>(())
93 /// ```
94 ///
95 /// # Errors
96 ///
97 /// Returns [`ValidationError::RankMismatch`] when the stride length does
98 /// not equal the compile-time rank.
99 fn strides_from_vec(strides: StrideVec) -> Result<Self::Strides>;
100
101 /// Convert this rank's stride representation into a dynamic stride vector.
102 ///
103 /// # Examples
104 ///
105 /// ```rust
106 /// use tenferro_tensor_core::{Rank, TensorRank};
107 ///
108 /// let strides = <Rank<2> as TensorRank>::strides_from_vec(vec![1, 2].into())?;
109 /// assert_eq!(<Rank<2> as TensorRank>::strides_into_vec(strides).as_slice(), &[1, 2]);
110 /// # Ok::<(), tenferro_tensor_core::ValidationError>(())
111 /// ```
112 fn strides_into_vec(strides: Self::Strides) -> StrideVec;
113}
114
115/// Dynamic tensor rank marker.
116///
117/// # Examples
118///
119/// ```rust
120/// use tenferro_tensor_core::{DynRank, TensorRank};
121///
122/// assert_eq!(DynRank::RANK, None);
123/// ```
124#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
125pub struct DynRank;
126
127/// Static tensor rank marker.
128///
129/// # Examples
130///
131/// ```rust
132/// use tenferro_tensor_core::{Rank, TensorRank};
133///
134/// assert_eq!(Rank::<3>::RANK, Some(3));
135/// ```
136#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
137pub struct Rank<const N: usize>;
138
139impl TensorRank for DynRank {
140 const RANK: Option<usize> = None;
141
142 type Shape = ShapeVec;
143 type Strides = StrideVec;
144
145 fn shape_from_vec(shape: ShapeVec) -> Result<Self::Shape> {
146 Ok(shape)
147 }
148
149 fn shape_into_vec(shape: Self::Shape) -> ShapeVec {
150 shape
151 }
152
153 fn strides_from_vec(strides: StrideVec) -> Result<Self::Strides> {
154 Ok(strides)
155 }
156
157 fn strides_into_vec(strides: Self::Strides) -> StrideVec {
158 strides
159 }
160}
161
162impl<const N: usize> TensorRank for Rank<N> {
163 const RANK: Option<usize> = Some(N);
164
165 type Shape = [usize; N];
166 type Strides = [isize; N];
167
168 fn shape_from_vec(shape: ShapeVec) -> Result<Self::Shape> {
169 let actual = shape.len();
170 shape
171 .into_vec()
172 .try_into()
173 .map_err(|_| ValidationError::RankMismatch {
174 expected: N,
175 actual,
176 })
177 }
178
179 fn shape_into_vec(shape: Self::Shape) -> ShapeVec {
180 ShapeVec::from_iter(shape)
181 }
182
183 fn strides_from_vec(strides: StrideVec) -> Result<Self::Strides> {
184 let actual = strides.len();
185 strides
186 .into_vec()
187 .try_into()
188 .map_err(|_| ValidationError::RankMismatch {
189 expected: N,
190 actual,
191 })
192 }
193
194 fn strides_into_vec(strides: Self::Strides) -> StrideVec {
195 StrideVec::from_iter(strides)
196 }
197}
198
199mod private {
200 pub trait Sealed {}
201
202 impl Sealed for super::DynRank {}
203 impl<const N: usize> Sealed for super::Rank<N> {}
204}