tenferro_ops/
sym_dim.rs

1use std::ops::{Add, Div, Mul, Sub};
2
3use thiserror::Error;
4
5use crate::dim_expr::DimExpr;
6
7/// A symbolic tensor dimension expression used to build shape-agnostic graphs.
8///
9/// `SymDim` values can be combined with basic arithmetic and later resolved
10/// against op-local [`DimExpr`] expressions when shape metadata is propagated.
11///
12/// # Examples
13///
14/// ```rust
15/// use tenferro_ops::sym_dim::SymDim;
16///
17/// let rows = SymDim::from(3usize);
18/// let cols = SymDim::from(4usize);
19/// let area = rows * cols;
20/// ```
21#[derive(Clone, Debug, PartialEq, Eq)]
22pub struct SymDim(pub(crate) RawSymDim);
23
24/// A symbolic dimension referenced a tensor that is not present in the
25/// conversion map supplied by the graph builder.
26#[derive(Clone, Debug, PartialEq, Eq, Error)]
27#[error("unknown symbolic tensor id {tensor_id}")]
28pub struct SymDimConversionError {
29    /// Identifier of the missing symbolic tensor.
30    pub tensor_id: u64,
31}
32
33#[derive(Clone, Debug, PartialEq, Eq)]
34pub(crate) enum RawSymDim {
35    Const(usize),
36    TensorAxis { tensor_id: u64, axis: usize },
37    Add(Box<RawSymDim>, Box<RawSymDim>),
38    Sub(Box<RawSymDim>, Box<RawSymDim>),
39    Mul(Box<RawSymDim>, Box<RawSymDim>),
40    FloorDiv(Box<RawSymDim>, Box<RawSymDim>),
41    Min(Box<RawSymDim>, Box<RawSymDim>),
42    Max(Box<RawSymDim>, Box<RawSymDim>),
43}
44
45impl SymDim {
46    /// Return the smaller of two symbolic dimensions.
47    ///
48    /// # Examples
49    ///
50    /// ```rust
51    /// use tenferro_ops::sym_dim::SymDim;
52    ///
53    /// let clipped = SymDim::from(5usize).min(8usize);
54    /// ```
55    pub fn min(self, other: impl Into<SymDim>) -> Self {
56        let other = other.into();
57        binary(self, other, |lhs, rhs| {
58            RawSymDim::Min(Box::new(lhs), Box::new(rhs))
59        })
60    }
61
62    /// Return the larger of two symbolic dimensions.
63    ///
64    /// # Examples
65    ///
66    /// ```rust
67    /// use tenferro_ops::sym_dim::SymDim;
68    ///
69    /// let batch = SymDim::from(2usize).max(1usize);
70    /// ```
71    pub fn max(self, other: impl Into<SymDim>) -> Self {
72        let other = other.into();
73        binary(self, other, |lhs, rhs| {
74            RawSymDim::Max(Box::new(lhs), Box::new(rhs))
75        })
76    }
77
78    #[doc(hidden)]
79    pub fn tensor_axis(tensor_id: u64, axis: usize) -> Self {
80        Self(RawSymDim::TensorAxis { tensor_id, axis })
81    }
82
83    #[doc(hidden)]
84    pub fn constant_value(&self) -> Option<usize> {
85        raw_constant_value(&self.0)
86    }
87
88    #[doc(hidden)]
89    pub fn from_dim_expr(expr: &DimExpr, input_shapes: &[&[SymDim]]) -> Self {
90        match expr {
91            DimExpr::Const(value) => Self::from(*value),
92            DimExpr::InputDim { input_idx, axis } => input_shapes[*input_idx][*axis].clone(),
93            DimExpr::Add(lhs, rhs) => {
94                Self::from_dim_expr(lhs, input_shapes) + Self::from_dim_expr(rhs, input_shapes)
95            }
96            DimExpr::Sub(lhs, rhs) => {
97                Self::from_dim_expr(lhs, input_shapes) - Self::from_dim_expr(rhs, input_shapes)
98            }
99            DimExpr::Mul(lhs, rhs) => {
100                Self::from_dim_expr(lhs, input_shapes) * Self::from_dim_expr(rhs, input_shapes)
101            }
102            DimExpr::FloorDiv(lhs, rhs) => {
103                Self::from_dim_expr(lhs, input_shapes) / Self::from_dim_expr(rhs, input_shapes)
104            }
105            DimExpr::Min(lhs, rhs) => {
106                Self::from_dim_expr(lhs, input_shapes).min(Self::from_dim_expr(rhs, input_shapes))
107            }
108            DimExpr::Max(lhs, rhs) => {
109                Self::from_dim_expr(lhs, input_shapes).max(Self::from_dim_expr(rhs, input_shapes))
110            }
111        }
112    }
113
114    #[doc(hidden)]
115    pub fn to_dim_expr(
116        &self,
117        tensor_map: &[(u64, usize)],
118    ) -> std::result::Result<DimExpr, SymDimConversionError> {
119        raw_to_dim_expr(&self.0, tensor_map)
120    }
121
122    /// Collect the unique traced tensor IDs referenced by `TensorAxis`
123    /// variants inside this expression, in traversal order (first
124    /// occurrence wins).
125    ///
126    /// Used by traced composition wrappers that build multi-input ops
127    /// from `SymDim`-valued target shapes — e.g.
128    /// `TracedTensor::broadcast_in_dim_sym`.
129    ///
130    /// # Examples
131    ///
132    /// ```rust
133    /// use tenferro_ops::sym_dim::SymDim;
134    ///
135    /// let lhs = SymDim::tensor_axis(7, 0);
136    /// let rhs = SymDim::tensor_axis(9, 1);
137    /// let sum = lhs + rhs;
138    /// assert_eq!(sum.referenced_tensor_ids(), vec![7, 9]);
139    /// ```
140    pub fn referenced_tensor_ids(&self) -> Vec<u64> {
141        let mut ids = Vec::new();
142        collect_tensor_ids(&self.0, &mut ids);
143        ids
144    }
145}
146
147fn collect_tensor_ids(raw: &RawSymDim, ids: &mut Vec<u64>) {
148    match raw {
149        RawSymDim::Const(_) => {}
150        RawSymDim::TensorAxis { tensor_id, .. } => {
151            if !ids.contains(tensor_id) {
152                ids.push(*tensor_id);
153            }
154        }
155        RawSymDim::Add(lhs, rhs)
156        | RawSymDim::Sub(lhs, rhs)
157        | RawSymDim::Mul(lhs, rhs)
158        | RawSymDim::FloorDiv(lhs, rhs)
159        | RawSymDim::Min(lhs, rhs)
160        | RawSymDim::Max(lhs, rhs) => {
161            collect_tensor_ids(lhs, ids);
162            collect_tensor_ids(rhs, ids);
163        }
164    }
165}
166
167fn raw_constant_value(raw: &RawSymDim) -> Option<usize> {
168    match raw {
169        RawSymDim::Const(value) => Some(*value),
170        RawSymDim::TensorAxis { .. } => None,
171        RawSymDim::Add(lhs, rhs) => raw_constant_value(lhs)?.checked_add(raw_constant_value(rhs)?),
172        RawSymDim::Sub(lhs, rhs) => raw_constant_value(lhs)?.checked_sub(raw_constant_value(rhs)?),
173        RawSymDim::Mul(lhs, rhs) => raw_constant_value(lhs)?.checked_mul(raw_constant_value(rhs)?),
174        RawSymDim::FloorDiv(lhs, rhs) => {
175            let rhs = raw_constant_value(rhs)?;
176            raw_constant_value(lhs)?.checked_div(rhs)
177        }
178        RawSymDim::Min(lhs, rhs) => Some(raw_constant_value(lhs)?.min(raw_constant_value(rhs)?)),
179        RawSymDim::Max(lhs, rhs) => Some(raw_constant_value(lhs)?.max(raw_constant_value(rhs)?)),
180    }
181}
182
183impl From<usize> for SymDim {
184    fn from(value: usize) -> Self {
185        Self(RawSymDim::Const(value))
186    }
187}
188
189fn raw_to_dim_expr(
190    raw: &RawSymDim,
191    tensor_map: &[(u64, usize)],
192) -> std::result::Result<DimExpr, SymDimConversionError> {
193    Ok(match raw {
194        RawSymDim::Const(value) => DimExpr::Const(*value),
195        RawSymDim::TensorAxis { tensor_id, axis } => {
196            let input_idx = tensor_map
197                .iter()
198                .find_map(|(candidate_id, input_idx)| {
199                    (candidate_id == tensor_id).then_some(*input_idx)
200                })
201                .ok_or(SymDimConversionError {
202                    tensor_id: *tensor_id,
203                })?;
204            DimExpr::InputDim {
205                input_idx,
206                axis: *axis,
207            }
208        }
209        RawSymDim::Add(lhs, rhs) => DimExpr::add(
210            raw_to_dim_expr(lhs, tensor_map)?,
211            raw_to_dim_expr(rhs, tensor_map)?,
212        ),
213        RawSymDim::Sub(lhs, rhs) => DimExpr::sub(
214            raw_to_dim_expr(lhs, tensor_map)?,
215            raw_to_dim_expr(rhs, tensor_map)?,
216        ),
217        RawSymDim::Mul(lhs, rhs) => DimExpr::mul(
218            raw_to_dim_expr(lhs, tensor_map)?,
219            raw_to_dim_expr(rhs, tensor_map)?,
220        ),
221        RawSymDim::FloorDiv(lhs, rhs) => DimExpr::floor_div(
222            raw_to_dim_expr(lhs, tensor_map)?,
223            raw_to_dim_expr(rhs, tensor_map)?,
224        ),
225        RawSymDim::Min(lhs, rhs) => DimExpr::min(
226            raw_to_dim_expr(lhs, tensor_map)?,
227            raw_to_dim_expr(rhs, tensor_map)?,
228        ),
229        RawSymDim::Max(lhs, rhs) => DimExpr::max(
230            raw_to_dim_expr(lhs, tensor_map)?,
231            raw_to_dim_expr(rhs, tensor_map)?,
232        ),
233    })
234}
235
236fn binary(lhs: SymDim, rhs: SymDim, f: impl FnOnce(RawSymDim, RawSymDim) -> RawSymDim) -> SymDim {
237    SymDim(f(lhs.0, rhs.0))
238}
239
240impl Add for SymDim {
241    type Output = SymDim;
242
243    fn add(self, rhs: SymDim) -> Self::Output {
244        binary(self, rhs, |lhs, rhs| {
245            RawSymDim::Add(Box::new(lhs), Box::new(rhs))
246        })
247    }
248}
249
250impl Add<usize> for SymDim {
251    type Output = SymDim;
252
253    fn add(self, rhs: usize) -> Self::Output {
254        self + SymDim::from(rhs)
255    }
256}
257
258impl Add<SymDim> for usize {
259    type Output = SymDim;
260
261    fn add(self, rhs: SymDim) -> Self::Output {
262        SymDim::from(self) + rhs
263    }
264}
265
266impl Sub for SymDim {
267    type Output = SymDim;
268
269    fn sub(self, rhs: SymDim) -> Self::Output {
270        binary(self, rhs, |lhs, rhs| {
271            RawSymDim::Sub(Box::new(lhs), Box::new(rhs))
272        })
273    }
274}
275
276impl Sub<usize> for SymDim {
277    type Output = SymDim;
278
279    fn sub(self, rhs: usize) -> Self::Output {
280        self - SymDim::from(rhs)
281    }
282}
283
284impl Sub<SymDim> for usize {
285    type Output = SymDim;
286
287    fn sub(self, rhs: SymDim) -> Self::Output {
288        SymDim::from(self) - rhs
289    }
290}
291
292impl Mul for SymDim {
293    type Output = SymDim;
294
295    fn mul(self, rhs: SymDim) -> Self::Output {
296        binary(self, rhs, |lhs, rhs| {
297            RawSymDim::Mul(Box::new(lhs), Box::new(rhs))
298        })
299    }
300}
301
302impl Mul<usize> for SymDim {
303    type Output = SymDim;
304
305    fn mul(self, rhs: usize) -> Self::Output {
306        self * SymDim::from(rhs)
307    }
308}
309
310impl Mul<SymDim> for usize {
311    type Output = SymDim;
312
313    fn mul(self, rhs: SymDim) -> Self::Output {
314        SymDim::from(self) * rhs
315    }
316}
317
318impl Div for SymDim {
319    type Output = SymDim;
320
321    fn div(self, rhs: SymDim) -> Self::Output {
322        binary(self, rhs, |lhs, rhs| {
323            RawSymDim::FloorDiv(Box::new(lhs), Box::new(rhs))
324        })
325    }
326}
327
328impl Div<usize> for SymDim {
329    type Output = SymDim;
330
331    fn div(self, rhs: usize) -> Self::Output {
332        self / SymDim::from(rhs)
333    }
334}
335
336impl Div<SymDim> for usize {
337    type Output = SymDim;
338
339    fn div(self, rhs: SymDim) -> Self::Output {
340        SymDim::from(self) / rhs
341    }
342}
tenferro_ops/sym_dim.rs

tenferro_ops/
sym_dim.rs
