tract_hir/infer/
factoid.rs

1use std::fmt;
2use std::iter::FromIterator;
3use std::ops::{Add, Div, Mul, Neg, Rem, Sub};
4
5use tract_num_traits::Zero;
6
7use crate::internal::*;
8
9/// Partial information about any value.
10pub trait Factoid: fmt::Debug + Clone + PartialEq + Default + Hash {
11    type Concrete: fmt::Debug;
12
13    /// Tries to transform the fact into a concrete value.
14    fn concretize(&self) -> Option<Self::Concrete>;
15
16    /// Returns whether the value is fully determined.
17    fn is_concrete(&self) -> bool {
18        self.concretize().is_some()
19    }
20
21    /// Tries to unify the fact with another fact of the same type.
22    fn unify(&self, other: &Self) -> TractResult<Self>;
23
24    /// Tries to unify the fact with another fact of the same type and update
25    /// self.
26    ///
27    /// Returns true if it actually changed something.
28    fn unify_with(&mut self, other: &Self) -> TractResult<bool> {
29        let new = self.unify(other)?;
30        let mut changed = false;
31        if &new != self {
32            changed = true;
33            *self = new;
34        }
35        Ok(changed)
36    }
37
38    /// Tries to unify the fact with another fact of the same type and update
39    /// both of them.
40    ///
41    /// Returns true if it actually changed something.
42    fn unify_with_mut(&mut self, other: &mut Self) -> TractResult<bool> {
43        let new = self.unify(other)?;
44        let mut changed = false;
45        if &new != self {
46            changed = true;
47            *self = new.clone();
48        }
49        if &new != other {
50            changed = true;
51            *other = new;
52        }
53        Ok(changed)
54    }
55
56    /// Tries to unify all facts in the list.
57    ///
58    ///
59    /// Returns true if it actually changed something.
60    fn unify_all(facts: &mut [&mut Self]) -> TractResult<bool> {
61        let mut overall_changed = false;
62        loop {
63            let mut changed = false;
64            for i in 0..facts.len() - 1 {
65                for j in i + 1..facts.len() {
66                    let (left, right) = facts.split_at_mut(j);
67                    let c = left[i].unify_with(right[0])?;
68                    changed = changed || c;
69                    overall_changed = changed || c;
70                }
71            }
72            if !changed {
73                return Ok(overall_changed);
74            }
75        }
76    }
77}
78
79/// Partial information about a value of type T.
80#[derive(Clone, PartialEq, Eq, Hash)]
81pub enum GenericFactoid<T: fmt::Debug + Clone + PartialEq + Hash> {
82    Only(T),
83    Any,
84}
85
86// if T is not Default, autoderive wont work
87#[allow(clippy::derivable_impls)]
88impl<T: fmt::Debug + Clone + PartialEq + Hash> Default for GenericFactoid<T> {
89    fn default() -> Self {
90        GenericFactoid::Any
91    }
92}
93
94impl<T: Copy + Clone + fmt::Debug + PartialEq + Hash> Copy for GenericFactoid<T> {}
95
96impl<T: fmt::Debug + Clone + PartialEq + Hash> Factoid for GenericFactoid<T> {
97    type Concrete = T;
98
99    /// Tries to transform the fact into a concrete value.
100    fn concretize(&self) -> Option<T> {
101        match self {
102            GenericFactoid::Any => None,
103            GenericFactoid::Only(m) => Some(m.clone()),
104        }
105    }
106
107    /// Tries to unify the fact with another fact of the same type.
108    fn unify(&self, other: &Self) -> TractResult<Self> {
109        let fact = match (self, other) {
110            (_, GenericFactoid::Any) => self.clone(),
111            (GenericFactoid::Any, _) => other.clone(),
112            _ if self == other => self.clone(),
113            _ => bail!("Impossible to unify {:?} with {:?}.", self, other),
114        };
115
116        Ok(fact)
117    }
118}
119
120impl<T: fmt::Debug + Clone + PartialEq + Hash> From<T> for GenericFactoid<T> {
121    fn from(t: T) -> Self {
122        GenericFactoid::Only(t)
123    }
124}
125
126impl<T: fmt::Display + fmt::Debug + Clone + PartialEq + Hash> fmt::Display for GenericFactoid<T> {
127    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
128        match self {
129            GenericFactoid::Any => write!(formatter, "?"),
130            GenericFactoid::Only(u) => write!(formatter, "{u}"),
131        }
132    }
133}
134
135impl<T: fmt::Debug + Clone + PartialEq + Hash> fmt::Debug for GenericFactoid<T> {
136    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
137        match self {
138            GenericFactoid::Any => write!(formatter, "?"),
139            GenericFactoid::Only(u) => write!(formatter, "{u:?}"),
140        }
141    }
142}
143
144/// Partial information about a type.
145pub type TypeFactoid = GenericFactoid<DatumType>;
146
147/// Partial information about a shape.
148///
149/// A basic example of a shape fact is `shapefactoid![1, 2]`, which corresponds to
150/// the shape `[1, 2]` in Arc<Tensor>. We can use `_` in facts to denote unknown
151/// dimensions (e.g. `shapefactoid![1, 2, _]` corresponds to any shape `[1, 2, k]`
152/// with `k` a non-negative integer). We can also use `..` at the end of a fact
153/// to only specify its first dimensions, so `shapefactoid![1, 2; ..]` matches any
154/// shape that starts with `[1, 2]` (e.g. `[1, 2, i]` or `[1, 2, i, j]`), while
155/// `shapefactoid![..]` matches any shape.
156#[derive(Clone, PartialEq, Eq, Hash)]
157pub struct ShapeFactoid {
158    pub(super) open: bool,
159    pub(super) dims: TVec<GenericFactoid<TDim>>,
160}
161
162impl ShapeFactoid {
163    /// Constructs an open shape fact.
164    pub fn open(dims: TVec<DimFact>) -> ShapeFactoid {
165        ShapeFactoid { open: true, dims }
166    }
167
168    pub fn is_open(&self) -> bool {
169        self.open
170    }
171
172    /// Constructs a closed shape fact.
173    pub fn closed(dims: TVec<DimFact>) -> ShapeFactoid {
174        ShapeFactoid { open: false, dims }
175    }
176
177    pub fn rank(&self) -> IntFactoid {
178        if self.open { GenericFactoid::Any } else { GenericFactoid::Only(self.dims.len() as i64) }
179    }
180
181    pub fn ensure_rank_at_least(&mut self, n: usize) -> bool {
182        let mut changed = false;
183        while self.dims.len() <= n {
184            self.dims.push(GenericFactoid::Any);
185            changed = true;
186        }
187        changed
188    }
189
190    pub fn dim(&self, i: usize) -> Option<DimFact> {
191        self.dims().nth(i).cloned()
192    }
193
194    pub fn set_dim(&mut self, i: usize, d: TDim) -> bool {
195        let fact = GenericFactoid::Only(d.clone());
196        if self.dim(i).as_ref() == Some(&fact) {
197            return false;
198        }
199        self.dims[i] = GenericFactoid::Only(d);
200        true
201    }
202
203    pub fn dims(&self) -> impl Iterator<Item = &DimFact> {
204        self.dims.iter()
205    }
206
207    pub fn as_concrete_finite(&self) -> TractResult<Option<TVec<usize>>> {
208        if self.open {
209            return Ok(None);
210        }
211        Ok(self.dims.iter().map(|d| d.concretize().and_then(|d| d.to_usize().ok())).collect())
212    }
213
214    pub fn matches(&self, t: &Tensor, symbols: Option<&SymbolValues>) -> TractResult<bool> {
215        let rank_compatible =
216            if self.is_open() { self.dims.len() <= t.rank() } else { self.dims.len() == t.rank() };
217        if !rank_compatible {
218            return Ok(false);
219        }
220
221        for i in 0..t.rank() {
222            let dim = self.dims.get(i).and_then(|el| el.concretize());
223            if let Some(dim) = dim.and_then(|dim| {
224                dim.eval(symbols.unwrap_or(&SymbolValues::default())).to_usize().ok()
225            }) {
226                if dim != t.shape()[i] {
227                    return Ok(false);
228                }
229            }
230        }
231        Ok(true)
232    }
233}
234
235impl Factoid for ShapeFactoid {
236    type Concrete = TVec<TDim>;
237
238    /// Tries to transform the fact into a `Vec<usize>`, or returns `None`.
239    fn concretize(self: &ShapeFactoid) -> Option<TVec<TDim>> {
240        if self.open {
241            return None;
242        }
243
244        let dims: TVec<_> = self.dims().filter_map(|d| d.concretize()).collect();
245
246        if dims.len() < self.dims.len() { None } else { Some(dims) }
247    }
248
249    /// Tries to unify the fact with another fact of the same type.
250    fn unify(&self, other: &Self) -> TractResult<Self> {
251        let (x, y) = (self, other);
252
253        use tract_itertools::EitherOrBoth::{Both, Left, Right};
254        use tract_itertools::Itertools;
255
256        let xi = x.dims();
257        let yi = y.dims();
258
259        let dimensions: TVec<_> = xi
260            .zip_longest(yi)
261            .map(|r| match r {
262                Both(a, b) => a.unify(b),
263                Left(d) if y.open => Ok(d.clone()),
264                Right(d) if x.open => Ok(d.clone()),
265
266                Left(_) | Right(_) => bail!(
267                    "Impossible to unify closed shapes of different rank (found {:?} and {:?}).",
268                    x,
269                    y
270                ),
271            })
272            .collect::<TractResult<_>>()
273            .with_context(|| format!("Unifying shapes {x:?} and {y:?}"))?;
274
275        if x.open && y.open {
276            Ok(ShapeFactoid::open(dimensions))
277        } else {
278            Ok(ShapeFactoid::closed(dimensions))
279        }
280    }
281}
282
283impl Default for ShapeFactoid {
284    /// Returns the most general shape fact possible.
285    fn default() -> ShapeFactoid {
286        ShapeFactoid::open(tvec![])
287    }
288}
289
290impl FromIterator<TDim> for ShapeFactoid {
291    /// Converts an iterator over usize into a closed shape.
292    fn from_iter<I: IntoIterator<Item = TDim>>(iter: I) -> ShapeFactoid {
293        ShapeFactoid::closed(iter.into_iter().map(GenericFactoid::Only).collect())
294    }
295}
296
297impl FromIterator<usize> for ShapeFactoid {
298    /// Converts an iterator over usize into a closed shape.
299    fn from_iter<I: IntoIterator<Item = usize>>(iter: I) -> ShapeFactoid {
300        ShapeFactoid::closed(iter.into_iter().map(|d| GenericFactoid::Only(d.to_dim())).collect())
301    }
302}
303
304impl<D: ToDim, I: IntoIterator<Item = D>> From<I> for ShapeFactoid {
305    fn from(it: I) -> ShapeFactoid {
306        ShapeFactoid::closed(it.into_iter().map(|d| GenericFactoid::Only(d.to_dim())).collect())
307    }
308}
309
310impl fmt::Debug for ShapeFactoid {
311    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
312        for (ix, d) in self.dims.iter().enumerate() {
313            if ix != 0 {
314                write!(formatter, ",")?
315            }
316            write!(formatter, "{d}")?;
317        }
318        if self.open {
319            if self.dims.len() == 0 {
320                write!(formatter, "..")?;
321            } else {
322                write!(formatter, ",..")?;
323            }
324        }
325        Ok(())
326    }
327}
328
329pub type DimFact = GenericFactoid<TDim>;
330
331/// Partial information about a value.
332pub type ValueFact = GenericFactoid<Arc<Tensor>>;
333
334pub type IntFactoid = GenericFactoid<i64>;
335
336impl<T> Zero for GenericFactoid<T>
337where
338    T: Add<T, Output = T> + Zero + PartialEq + Clone + ::std::fmt::Debug + Hash,
339{
340    fn zero() -> GenericFactoid<T> {
341        GenericFactoid::Only(T::zero())
342    }
343    fn is_zero(&self) -> bool {
344        match self {
345            GenericFactoid::Only(t) => t.is_zero(),
346            _ => false,
347        }
348    }
349}
350
351impl<T> Neg for GenericFactoid<T>
352where
353    T: Neg<Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
354{
355    type Output = GenericFactoid<T>;
356    fn neg(self) -> GenericFactoid<T> {
357        match self {
358            GenericFactoid::Only(t) => GenericFactoid::Only(t.neg()),
359            any => any,
360        }
361    }
362}
363
364impl<T, I> Add<I> for GenericFactoid<T>
365where
366    T: Add<T, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
367    I: Into<GenericFactoid<T>>,
368{
369    type Output = GenericFactoid<T>;
370    fn add(self, rhs: I) -> Self::Output {
371        match (self.concretize(), rhs.into().concretize()) {
372            (Some(a), Some(b)) => GenericFactoid::Only(a + b),
373            _ => GenericFactoid::Any,
374        }
375    }
376}
377
378impl<T> Sub<GenericFactoid<T>> for GenericFactoid<T>
379where
380    T: Sub<T, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
381{
382    type Output = GenericFactoid<T>;
383    fn sub(self, rhs: GenericFactoid<T>) -> Self::Output {
384        match (self.concretize(), rhs.concretize()) {
385            (Some(a), Some(b)) => GenericFactoid::Only(a - b),
386            _ => GenericFactoid::Any,
387        }
388    }
389}
390
391impl<T, R> Mul<R> for GenericFactoid<T>
392where
393    T: Mul<R, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
394{
395    type Output = GenericFactoid<T>;
396    fn mul(self, rhs: R) -> Self::Output {
397        if let Some(a) = self.concretize() {
398            GenericFactoid::Only(a * rhs)
399        } else {
400            GenericFactoid::Any
401        }
402    }
403}
404
405impl<T, R> Div<R> for GenericFactoid<T>
406where
407    T: Div<R, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
408{
409    type Output = GenericFactoid<T>;
410    fn div(self, rhs: R) -> Self::Output {
411        if let Some(a) = self.concretize() {
412            GenericFactoid::Only(a / rhs)
413        } else {
414            GenericFactoid::Any
415        }
416    }
417}
418
419impl<T, R> Rem<R> for GenericFactoid<T>
420where
421    T: Rem<R, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
422{
423    type Output = GenericFactoid<T>;
424    fn rem(self, rhs: R) -> Self::Output {
425        if let Some(a) = self.concretize() {
426            GenericFactoid::Only(a % rhs)
427        } else {
428            GenericFactoid::Any
429        }
430    }
431}
432
433#[cfg(test)]
434mod tests {
435    use super::GenericFactoid::*;
436    use super::*;
437
438    #[test]
439    fn unify_same_datum_type() {
440        let dt = TypeFactoid::Only(DatumType::F32);
441        assert_eq!(dt.unify(&dt).unwrap(), dt);
442    }
443
444    #[test]
445    fn unify_different_datum_types_only() {
446        let dt1 = TypeFactoid::Only(DatumType::F32);
447        let dt2 = TypeFactoid::Only(DatumType::F64);
448        assert!(dt1.unify(&dt2).is_err());
449    }
450
451    #[test]
452    fn unify_different_datum_types_any_left() {
453        let dt = TypeFactoid::Only(DatumType::F32);
454        assert_eq!(TypeFactoid::Any.unify(&dt).unwrap(), dt);
455    }
456
457    #[test]
458    fn unify_different_datum_types_any_right() {
459        let dt = TypeFactoid::Only(DatumType::F32);
460        assert_eq!(dt.unify(&TypeFactoid::Any).unwrap(), dt);
461    }
462
463    #[test]
464    fn unify_same_shape_1() {
465        let s = ShapeFactoid::closed(tvec![]);
466        assert_eq!(s.unify(&s).unwrap(), s);
467    }
468
469    #[test]
470    fn unify_same_shape_2() {
471        let s = ShapeFactoid::closed(tvec![Any]);
472        assert_eq!(s.unify(&s).unwrap(), s);
473    }
474
475    #[test]
476    fn unify_same_shape_3() {
477        let s = ShapeFactoid::closed(tvec![Only(1.into()), Only(2.into())]);
478        assert_eq!(s.unify(&s).unwrap(), s);
479    }
480
481    #[test]
482    fn unify_different_shapes_1() {
483        let s1 = ShapeFactoid::closed(tvec![Only(1.into()), Only(2.into())]);
484        let s2 = ShapeFactoid::closed(tvec![Only(1.into())]);
485        assert!(s1.unify(&s2).is_err());
486    }
487
488    #[test]
489    fn unify_different_shapes_2() {
490        let s1 = ShapeFactoid::closed(tvec![Only(1.into()), Only(2.into())]);
491        let s2 = ShapeFactoid::closed(tvec![Any]);
492        assert!(s1.unify(&s2).is_err());
493    }
494
495    #[test]
496    fn unify_different_shapes_3() {
497        let s1 = ShapeFactoid::open(tvec![Only(1.into()), Only(2.into())]);
498        let s2 = ShapeFactoid::closed(tvec![Any]);
499        assert!(s1.unify(&s2).is_err());
500    }
501
502    #[test]
503    fn unify_different_shapes_4() {
504        let s1 = ShapeFactoid::closed(tvec![Any]);
505        let s2 = ShapeFactoid::closed(tvec![Any]);
506        let sr = ShapeFactoid::closed(tvec![Any]);
507        assert_eq!(s1.unify(&s2).unwrap(), sr);
508    }
509
510    #[test]
511    fn unify_different_shapes_5() {
512        let s1 = ShapeFactoid::closed(tvec![Any]);
513        let s2 = ShapeFactoid::closed(tvec![Only(1.into())]);
514        let sr = ShapeFactoid::closed(tvec![Only(1.into())]);
515        assert_eq!(s1.unify(&s2).unwrap(), sr);
516    }
517
518    #[test]
519    fn unify_different_shapes_6() {
520        let s1 = ShapeFactoid::open(tvec![]);
521        let s2 = ShapeFactoid::closed(tvec![Only(1.into())]);
522        let sr = ShapeFactoid::closed(tvec![Only(1.into())]);
523        assert_eq!(s1.unify(&s2).unwrap(), sr);
524    }
525
526    #[test]
527    fn unify_different_shapes_7() {
528        let s1 = ShapeFactoid::open(tvec![Any, Only(2.into())]);
529        let s2 = ShapeFactoid::closed(tvec![Only(1.into()), Any, Any]);
530        let sr = ShapeFactoid::closed(tvec![Only(1.into()), Only(2.into()), Any]);
531        assert_eq!(s1.unify(&s2).unwrap(), sr);
532    }
533
534    #[test]
535    fn unify_same_value() {
536        let t = ValueFact::Only(rctensor0(12f32));
537        assert_eq!(t.unify(&t).unwrap(), t);
538    }
539
540    #[test]
541    fn unify_different_values_only() {
542        let t1 = ValueFact::Only(rctensor1(&[12f32]));
543        let t2 = ValueFact::Only(rctensor1(&[12f32, 42.0]));
544        assert!(t1.unify(&t2).is_err());
545    }
546
547    #[test]
548    fn unify_different_values_any_left() {
549        let t1 = ValueFact::Only(rctensor1(&[12f32]));
550        assert_eq!(ValueFact::Any.unify(&t1).unwrap(), t1);
551    }
552
553    #[test]
554    fn unify_different_values_any_right() {
555        let t1 = ValueFact::Only(rctensor1(&[12f32]));
556        assert_eq!(t1.unify(&ValueFact::Any).unwrap(), t1);
557    }
558}
tract_hir/infer/factoid.rs

tract_hir/infer/
factoid.rs
