Skip to main content

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}