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        rule_if!(!self.open);
209        Ok(self.dims.iter().map(|d| d.concretize().and_then(|d| d.to_usize().ok())).collect())
210    }
211
212    pub fn matches(&self, t: &Tensor, symbols: Option<&SymbolValues>) -> TractResult<bool> {
213        let rank_compatible =
214            if self.is_open() { self.dims.len() <= t.rank() } else { self.dims.len() == t.rank() };
215        if !rank_compatible {
216            return Ok(false);
217        }
218
219        for i in 0..t.rank() {
220            let dim = self.dims.get(i).and_then(|el| el.concretize());
221            if let Some(dim) = dim.and_then(|dim| {
222                dim.eval(symbols.unwrap_or(&SymbolValues::default())).to_usize().ok()
223            }) {
224                if dim != t.shape()[i] {
225                    return Ok(false);
226                }
227            }
228        }
229        Ok(true)
230    }
231}
232
233impl Factoid for ShapeFactoid {
234    type Concrete = TVec<TDim>;
235
236    /// Tries to transform the fact into a `Vec<usize>`, or returns `None`.
237    fn concretize(self: &ShapeFactoid) -> Option<TVec<TDim>> {
238        if self.open {
239            return None;
240        }
241
242        let dims: TVec<_> = self.dims().filter_map(|d| d.concretize()).collect();
243
244        if dims.len() < self.dims.len() { None } else { Some(dims) }
245    }
246
247    /// Tries to unify the fact with another fact of the same type.
248    fn unify(&self, other: &Self) -> TractResult<Self> {
249        let (x, y) = (self, other);
250
251        use tract_itertools::EitherOrBoth::{Both, Left, Right};
252        use tract_itertools::Itertools;
253
254        let xi = x.dims();
255        let yi = y.dims();
256
257        let dimensions: TVec<_> = xi
258            .zip_longest(yi)
259            .map(|r| match r {
260                Both(a, b) => a.unify(b),
261                Left(d) if y.open => Ok(d.clone()),
262                Right(d) if x.open => Ok(d.clone()),
263
264                Left(_) | Right(_) => bail!(
265                    "Impossible to unify closed shapes of different rank (found {:?} and {:?}).",
266                    x,
267                    y
268                ),
269            })
270            .collect::<TractResult<_>>()
271            .with_context(|| format!("Unifying shapes {x:?} and {y:?}"))?;
272
273        if x.open && y.open {
274            Ok(ShapeFactoid::open(dimensions))
275        } else {
276            Ok(ShapeFactoid::closed(dimensions))
277        }
278    }
279}
280
281impl Default for ShapeFactoid {
282    /// Returns the most general shape fact possible.
283    fn default() -> ShapeFactoid {
284        ShapeFactoid::open(tvec![])
285    }
286}
287
288impl FromIterator<TDim> for ShapeFactoid {
289    /// Converts an iterator over usize into a closed shape.
290    fn from_iter<I: IntoIterator<Item = TDim>>(iter: I) -> ShapeFactoid {
291        ShapeFactoid::closed(iter.into_iter().map(GenericFactoid::Only).collect())
292    }
293}
294
295impl FromIterator<usize> for ShapeFactoid {
296    /// Converts an iterator over usize into a closed shape.
297    fn from_iter<I: IntoIterator<Item = usize>>(iter: I) -> ShapeFactoid {
298        ShapeFactoid::closed(iter.into_iter().map(|d| GenericFactoid::Only(d.to_dim())).collect())
299    }
300}
301
302impl<D: ToDim, I: IntoIterator<Item = D>> From<I> for ShapeFactoid {
303    fn from(it: I) -> ShapeFactoid {
304        ShapeFactoid::closed(it.into_iter().map(|d| GenericFactoid::Only(d.to_dim())).collect())
305    }
306}
307
308impl fmt::Debug for ShapeFactoid {
309    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
310        for (ix, d) in self.dims.iter().enumerate() {
311            if ix != 0 {
312                write!(formatter, ",")?
313            }
314            write!(formatter, "{d}")?;
315        }
316        if self.open {
317            if self.dims.len() == 0 {
318                write!(formatter, "..")?;
319            } else {
320                write!(formatter, ",..")?;
321            }
322        }
323        Ok(())
324    }
325}
326
327pub type DimFact = GenericFactoid<TDim>;
328
329/// Partial information about a value.
330pub type ValueFact = GenericFactoid<Arc<Tensor>>;
331
332pub type IntFactoid = GenericFactoid<i64>;
333
334impl<T> Zero for GenericFactoid<T>
335where
336    T: Add<T, Output = T> + Zero + PartialEq + Clone + ::std::fmt::Debug + Hash,
337{
338    fn zero() -> GenericFactoid<T> {
339        GenericFactoid::Only(T::zero())
340    }
341    fn is_zero(&self) -> bool {
342        match self {
343            GenericFactoid::Only(t) => t.is_zero(),
344            _ => false,
345        }
346    }
347}
348
349impl<T> Neg for GenericFactoid<T>
350where
351    T: Neg<Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
352{
353    type Output = GenericFactoid<T>;
354    fn neg(self) -> GenericFactoid<T> {
355        match self {
356            GenericFactoid::Only(t) => GenericFactoid::Only(t.neg()),
357            any => any,
358        }
359    }
360}
361
362impl<T, I> Add<I> for GenericFactoid<T>
363where
364    T: Add<T, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
365    I: Into<GenericFactoid<T>>,
366{
367    type Output = GenericFactoid<T>;
368    fn add(self, rhs: I) -> Self::Output {
369        match (self.concretize(), rhs.into().concretize()) {
370            (Some(a), Some(b)) => GenericFactoid::Only(a + b),
371            _ => GenericFactoid::Any,
372        }
373    }
374}
375
376impl<T> Sub<GenericFactoid<T>> for GenericFactoid<T>
377where
378    T: Sub<T, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
379{
380    type Output = GenericFactoid<T>;
381    fn sub(self, rhs: GenericFactoid<T>) -> Self::Output {
382        match (self.concretize(), rhs.concretize()) {
383            (Some(a), Some(b)) => GenericFactoid::Only(a - b),
384            _ => GenericFactoid::Any,
385        }
386    }
387}
388
389impl<T, R> Mul<R> for GenericFactoid<T>
390where
391    T: Mul<R, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
392{
393    type Output = GenericFactoid<T>;
394    fn mul(self, rhs: R) -> Self::Output {
395        if let Some(a) = self.concretize() {
396            GenericFactoid::Only(a * rhs)
397        } else {
398            GenericFactoid::Any
399        }
400    }
401}
402
403impl<T, R> Div<R> for GenericFactoid<T>
404where
405    T: Div<R, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
406{
407    type Output = GenericFactoid<T>;
408    fn div(self, rhs: R) -> Self::Output {
409        if let Some(a) = self.concretize() {
410            GenericFactoid::Only(a / rhs)
411        } else {
412            GenericFactoid::Any
413        }
414    }
415}
416
417impl<T, R> Rem<R> for GenericFactoid<T>
418where
419    T: Rem<R, Output = T> + PartialEq + Clone + ::std::fmt::Debug + Hash,
420{
421    type Output = GenericFactoid<T>;
422    fn rem(self, rhs: R) -> Self::Output {
423        if let Some(a) = self.concretize() {
424            GenericFactoid::Only(a % rhs)
425        } else {
426            GenericFactoid::Any
427        }
428    }
429}
430
431#[cfg(test)]
432mod tests {
433    use super::GenericFactoid::*;
434    use super::*;
435
436    #[test]
437    fn unify_same_datum_type() {
438        let dt = TypeFactoid::Only(DatumType::F32);
439        assert_eq!(dt.unify(&dt).unwrap(), dt);
440    }
441
442    #[test]
443    fn unify_different_datum_types_only() {
444        let dt1 = TypeFactoid::Only(DatumType::F32);
445        let dt2 = TypeFactoid::Only(DatumType::F64);
446        assert!(dt1.unify(&dt2).is_err());
447    }
448
449    #[test]
450    fn unify_different_datum_types_any_left() {
451        let dt = TypeFactoid::Only(DatumType::F32);
452        assert_eq!(TypeFactoid::Any.unify(&dt).unwrap(), dt);
453    }
454
455    #[test]
456    fn unify_different_datum_types_any_right() {
457        let dt = TypeFactoid::Only(DatumType::F32);
458        assert_eq!(dt.unify(&TypeFactoid::Any).unwrap(), dt);
459    }
460
461    #[test]
462    fn unify_same_shape_1() {
463        let s = ShapeFactoid::closed(tvec![]);
464        assert_eq!(s.unify(&s).unwrap(), s);
465    }
466
467    #[test]
468    fn unify_same_shape_2() {
469        let s = ShapeFactoid::closed(tvec![Any]);
470        assert_eq!(s.unify(&s).unwrap(), s);
471    }
472
473    #[test]
474    fn unify_same_shape_3() {
475        let s = ShapeFactoid::closed(tvec![Only(1.into()), Only(2.into())]);
476        assert_eq!(s.unify(&s).unwrap(), s);
477    }
478
479    #[test]
480    fn unify_different_shapes_1() {
481        let s1 = ShapeFactoid::closed(tvec![Only(1.into()), Only(2.into())]);
482        let s2 = ShapeFactoid::closed(tvec![Only(1.into())]);
483        assert!(s1.unify(&s2).is_err());
484    }
485
486    #[test]
487    fn unify_different_shapes_2() {
488        let s1 = ShapeFactoid::closed(tvec![Only(1.into()), Only(2.into())]);
489        let s2 = ShapeFactoid::closed(tvec![Any]);
490        assert!(s1.unify(&s2).is_err());
491    }
492
493    #[test]
494    fn unify_different_shapes_3() {
495        let s1 = ShapeFactoid::open(tvec![Only(1.into()), Only(2.into())]);
496        let s2 = ShapeFactoid::closed(tvec![Any]);
497        assert!(s1.unify(&s2).is_err());
498    }
499
500    #[test]
501    fn unify_different_shapes_4() {
502        let s1 = ShapeFactoid::closed(tvec![Any]);
503        let s2 = ShapeFactoid::closed(tvec![Any]);
504        let sr = ShapeFactoid::closed(tvec![Any]);
505        assert_eq!(s1.unify(&s2).unwrap(), sr);
506    }
507
508    #[test]
509    fn unify_different_shapes_5() {
510        let s1 = ShapeFactoid::closed(tvec![Any]);
511        let s2 = ShapeFactoid::closed(tvec![Only(1.into())]);
512        let sr = ShapeFactoid::closed(tvec![Only(1.into())]);
513        assert_eq!(s1.unify(&s2).unwrap(), sr);
514    }
515
516    #[test]
517    fn unify_different_shapes_6() {
518        let s1 = ShapeFactoid::open(tvec![]);
519        let s2 = ShapeFactoid::closed(tvec![Only(1.into())]);
520        let sr = ShapeFactoid::closed(tvec![Only(1.into())]);
521        assert_eq!(s1.unify(&s2).unwrap(), sr);
522    }
523
524    #[test]
525    fn unify_different_shapes_7() {
526        let s1 = ShapeFactoid::open(tvec![Any, Only(2.into())]);
527        let s2 = ShapeFactoid::closed(tvec![Only(1.into()), Any, Any]);
528        let sr = ShapeFactoid::closed(tvec![Only(1.into()), Only(2.into()), Any]);
529        assert_eq!(s1.unify(&s2).unwrap(), sr);
530    }
531
532    #[test]
533    fn unify_same_value() {
534        let t = ValueFact::Only(rctensor0(12f32));
535        assert_eq!(t.unify(&t).unwrap(), t);
536    }
537
538    #[test]
539    fn unify_different_values_only() {
540        let t1 = ValueFact::Only(rctensor1(&[12f32]));
541        let t2 = ValueFact::Only(rctensor1(&[12f32, 42.0]));
542        assert!(t1.unify(&t2).is_err());
543    }
544
545    #[test]
546    fn unify_different_values_any_left() {
547        let t1 = ValueFact::Only(rctensor1(&[12f32]));
548        assert_eq!(ValueFact::Any.unify(&t1).unwrap(), t1);
549    }
550
551    #[test]
552    fn unify_different_values_any_right() {
553        let t1 = ValueFact::Only(rctensor1(&[12f32]));
554        assert_eq!(t1.unify(&ValueFact::Any).unwrap(), t1);
555    }
556}
tract_hir/infer/factoid.rs

tract_hir/infer/
factoid.rs
