tract_data/dim/
mod.rs

1//! Extended dimension support
2use crate::internal::*;
3use num_traits::{One, Zero};
4use std::fmt;
5use std::ops;
6
7mod assertion;
8mod parse;
9mod resolve;
10mod sym;
11mod tree;
12
13pub use self::assertion::Assertion;
14pub use self::parse::parse_tdim;
15pub use self::resolve::solve_for;
16pub use self::sym::{Symbol, SymbolScope, SymbolValues};
17pub use self::tree::{TDim, TooEarly};
18
19use crate::{TractError, TractResult};
20
21/// A super-trait for value acting as tensor dimensions in tract.
22///
23/// Implemented by:
24///
25/// * `usize` for regular dimensions
26/// * `TDim` supporting regular and streaming dimensions
27pub trait DimLike:
28    Clone
29    + Default
30    + PartialEq
31    + From<usize>
32    + for<'a> std::convert::TryFrom<&'a TDim, Error = TractError>
33    + ::num_traits::Zero
34    + fmt::Debug
35    + fmt::Display
36    + std::hash::Hash
37    + ops::Add<Self, Output = Self>
38    + ops::Add<usize, Output = Self>
39    + for<'a> ops::Add<&'a Self, Output = Self>
40    + ops::Sub<Self, Output = Self>
41    + ops::Sub<usize, Output = Self>
42    + for<'a> ops::Sub<&'a Self, Output = Self>
43    + ops::Mul<Self, Output = Self>
44    + ops::Mul<usize, Output = Self>
45    + for<'a> ops::Mul<&'a Self, Output = Self>
46    + ops::Div<usize, Output = Self>
47    + ops::Rem<usize, Output = Self>
48    + Send
49    + Sync
50    + 'static
51    + std::iter::Sum
52    + std::iter::Product
53    + ToDim
54    + One
55{
56    fn maybe_div(&self, other: &Self) -> TractResult<(Self, u64)>;
57
58    /// Integer divise, rounding up to next integer.
59    fn divceil(&self, other: usize) -> Self {
60        (self.clone() + other - 1) / other
61    }
62
63    /// Convert to regular integer.
64    fn to_i64(&self) -> TractResult<i64>;
65
66    fn to_usize(&self) -> TractResult<usize> {
67        self.to_i64().map(|d| d as usize)
68    }
69
70    fn to_isize(&self) -> TractResult<isize> {
71        self.to_i64().map(|d| d as isize)
72    }
73
74    fn to_i32(&self) -> TractResult<i32> {
75        self.to_i64().map(|d| d as i32)
76    }
77
78    /// Substitute as many symbols as possible in the dim value.
79    fn eval(&self, values: &SymbolValues) -> Self;
80
81    /// Full evaluation of the symbol, failing if a symbol is missing
82    fn eval_to_i64(&self, values: &SymbolValues) -> TractResult<i64>;
83
84    fn substitute(&self, from: &Symbol, to: &Self) -> TractResult<Self>;
85    fn substitute_all(&self, map: &std::collections::HashMap<Symbol, Self>) -> TractResult<Self>;
86
87    fn broadcast(self, other: Self) -> TractResult<Self>;
88    fn mini(self, other: Self) -> Self;
89    fn maxi(self, other: Self) -> Self;
90
91    fn compatible_with(&self, other: &Self) -> bool;
92}
93
94impl DimLike for TDim {
95    fn maybe_div(&self, other: &Self) -> TractResult<(Self, u64)> {
96        if self.is_zero() {
97            return Ok((TDim::zero(), 1));
98        } else if other.is_zero() {
99            bail!("Division by zero")
100        }
101        fn expand(dim: &TDim) -> (i64, Vec<TDim>) {
102            match dim {
103                TDim::Mul(terms) => terms.iter().map(expand).fold((1i64, vec![]), |acc, t| {
104                    (acc.0 * t.0, acc.1.into_iter().chain(t.1).collect())
105                }),
106                TDim::MulInt(a, terms) => {
107                    let (b, v) = expand(terms);
108                    (a * b, v)
109                }
110                TDim::Val(x) => (*x, vec![]),
111                TDim::Add(terms) => {
112                    let gcd =
113                        terms.iter().map(expand).map(|(n, _)| n).reduce(|a, b| a.gcd(&b)).unwrap();
114                    (
115                        gcd,
116                        vec![TDim::Add(terms.iter().map(|t| t.clone() / gcd).collect()).simplify()],
117                    )
118                }
119                it => (1, vec![it.clone()]),
120            }
121        }
122        let (mut num_int, mut num) = expand(self);
123        let (mut denum_int, mut denum) = expand(other);
124        if num == denum {
125            num = vec![];
126            denum = vec![];
127        }
128        for it in denum {
129            if let Some(pos) = num.iter().position(|n| n == &it) {
130                num.remove(pos);
131            } else {
132                bail!("Can't divide {} by {}", self, other)
133            }
134        }
135        use num_integer::Integer;
136        if denum_int < 0 {
137            num_int *= -1;
138            denum_int *= -1;
139        }
140        let gcd = num_int.gcd(&denum_int);
141        num_int /= gcd;
142        denum_int /= gcd;
143        Ok(((TDim::Mul(num) * num_int).reduce(), denum_int as u64))
144    }
145
146    fn to_i64(&self) -> TractResult<i64> {
147        TDim::to_i64(self)
148    }
149
150    fn eval(&self, values: &SymbolValues) -> Self {
151        self.eval(values)
152    }
153
154    fn substitute(&self, from: &Symbol, to: &Self) -> TractResult<Self> {
155        self.substitute(from, to)
156    }
157
158    fn substitute_all(&self, map: &std::collections::HashMap<Symbol, Self>) -> TractResult<Self> {
159        TDim::substitute_all(self, map)
160    }
161
162    fn eval_to_i64(&self, values: &SymbolValues) -> TractResult<i64> {
163        TDim::eval_to_i64(self, values)
164    }
165
166    fn broadcast(self, other: Self) -> TractResult<Self> {
167        if self.is_one() {
168            Ok(other)
169        } else if other.is_one() {
170            Ok(self)
171        } else {
172            Ok(TDim::Broadcast(vec![self, other]).simplify())
173        }
174    }
175
176    fn compatible_with(&self, other: &Self) -> bool {
177        self.compatible_with(other)
178    }
179
180    fn mini(self, other: Self) -> Self {
181        TDim::Min(vec![self, other]).simplify()
182    }
183
184    fn maxi(self, other: Self) -> Self {
185        TDim::Max(vec![self, other]).simplify()
186    }
187}
188
189impl<'a> std::convert::TryFrom<&'a TDim> for TDim {
190    type Error = TractError;
191    fn try_from(d: &'a TDim) -> TractResult<TDim> {
192        Ok(d.clone())
193    }
194}
195
196impl DimLike for usize {
197    fn maybe_div(&self, other: &Self) -> TractResult<(Self, u64)> {
198        use num_integer::Integer;
199        let gcd = self.gcd(other);
200        Ok((self / gcd, (other / gcd) as u64))
201    }
202
203    fn to_i64(&self) -> TractResult<i64> {
204        Ok(*self as i64)
205    }
206
207    fn eval(&self, _values: &SymbolValues) -> Self {
208        *self
209    }
210
211    fn substitute(&self, _from: &Symbol, _to: &Self) -> TractResult<Self> {
212        Ok(*self)
213    }
214
215    fn substitute_all(&self, _map: &std::collections::HashMap<Symbol, Self>) -> TractResult<Self> {
216        Ok(*self)
217    }
218
219    fn eval_to_i64(&self, _: &SymbolValues) -> TractResult<i64> {
220        Ok(*self as i64)
221    }
222
223    fn broadcast(self, other: Self) -> TractResult<Self> {
224        if self == 1 || self == other {
225            Ok(other)
226        } else if other == 1 {
227            Ok(self)
228        } else {
229            bail!("Can not broadcast {self} against {other}")
230        }
231    }
232
233    fn compatible_with(&self, other: &Self) -> bool {
234        self == other
235    }
236
237    fn mini(self, other: Self) -> Self {
238        if self < other { self } else { other }
239    }
240
241    fn maxi(self, other: Self) -> Self {
242        if self > other { self } else { other }
243    }
244}
245
246impl<'a> std::convert::TryFrom<&'a TDim> for usize {
247    type Error = TractError;
248    fn try_from(d: &'a TDim) -> TractResult<usize> {
249        d.to_usize()
250    }
251}
252
253/// Convenience trait to convert values to TDim.
254pub trait ToDim {
255    /// Convert self to a TDim.
256    fn to_dim(&self) -> TDim;
257}
258
259impl<I: Into<TDim> + Clone> ToDim for I {
260    fn to_dim(&self) -> TDim {
261        self.clone().into()
262    }
263}
264
265impl ToDim for &TDim {
266    fn to_dim(&self) -> TDim {
267        (*self).clone()
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    lazy_static::lazy_static! {
276        static ref S: (SymbolScope, Symbol) = {
277            let table = SymbolScope::default();
278            let s = table.new_with_prefix("S");
279            (table, s)
280        };
281    }
282
283    pub fn s() -> TDim {
284        S.1.clone().into()
285    }
286
287    #[test]
288    fn div() {
289        assert_eq!(TDim::from(12).maybe_div(&TDim::from(4)).unwrap(), (3.into(), 1));
290    }
291
292    #[test]
293    fn div_sym_int() {
294        assert_eq!((s() * 12).maybe_div(&TDim::from(4)).unwrap(), (s() * 3, 1));
295    }
296
297    #[test]
298    fn div_sym_sym() {
299        assert_eq!((s() * 12).maybe_div(&(s() * 4)).unwrap(), (3.into(), 1));
300    }
301
302    #[test]
303    fn div_sym_sym_ratio() {
304        assert_eq!((s() * 13).maybe_div(&(s() * 4)).unwrap(), (13.into(), 4));
305    }
306
307    #[test]
308    fn div_sym_sym_rem() {
309        assert!((s() + 1).maybe_div(&(s() * 4)).is_err());
310    }
311
312    #[test]
313    fn div_sym_sym_simply_1() {
314        assert_eq!((s()).maybe_div(&(s())).unwrap(), (TDim::Val(1), 1));
315    }
316
317    #[test]
318    fn div_sym_sym_complex() {
319        let s = s();
320        let b = S.0.sym("b");
321        assert_eq!(
322            (256.to_dim() * &s * &b).maybe_div(&(1.to_dim() * &s * &b)).unwrap(),
323            (256.into(), 1)
324        );
325    }
326
327    #[test]
328    fn div_sym_sym_with_add() {
329        assert_eq!((s() * 80 - 160).maybe_div(&(s() - 2)).unwrap(), (80.into(), 1));
330    }
331}
tract_data/dim/mod.rs

tract_data/dim/
mod.rs
