tract_hir/infer/
fact.rs

1use std::convert::TryFrom;
2use std::fmt;
3use std::sync::Arc;
4
5use super::factoid::*;
6use crate::internal::*;
7
8/// Partial information about a tensor.
9///
10/// The task of the analyser is to tag every edge in the graph with information
11/// about the tensors that flow through it - specifically their datum_type, their
12/// shape and possibly their value. During the analysis, however, we might only
13/// know some of that information (say, for instance, that an edge only carries
14/// tensors of rank 4, but without knowing their precise dimension).
15///
16/// This is where tensor facts come in: they hold partial information about the
17/// datum_type, shape and value of tensors that might flow through an edge of the
18/// graph. The analyser will first tag each edge with a fact, starting with the
19/// most general one and specializing it at each iteration. Eventually, it will
20/// reach a fixed point that - hopefully - holds enough information.
21#[derive(Clone, PartialEq, Eq, Default, Hash)]
22pub struct InferenceFact {
23    pub datum_type: TypeFactoid,
24    pub shape: ShapeFactoid,
25    pub value: ValueFact,
26}
27
28impl InferenceFact {
29    /// Constructs the most general tensor fact possible.
30    pub fn new() -> InferenceFact {
31        InferenceFact::default()
32    }
33
34    pub fn any() -> InferenceFact {
35        InferenceFact::default()
36    }
37
38    pub fn dt(dt: DatumType) -> InferenceFact {
39        InferenceFact::default().with_datum_type(dt)
40    }
41
42    pub fn dt_shape<S: Into<ShapeFactoid>>(dt: DatumType, shape: S) -> InferenceFact {
43        InferenceFact::dt(dt).with_shape(shape)
44    }
45
46    pub fn shape<S: Into<ShapeFactoid>>(shape: S) -> InferenceFact {
47        InferenceFact::default().with_shape(shape)
48    }
49
50    pub fn with_datum_type(self, dt: DatumType) -> InferenceFact {
51        InferenceFact { datum_type: dt.into(), ..self }
52    }
53
54    pub fn without_datum_type(self) -> InferenceFact {
55        InferenceFact { datum_type: TypeFactoid::Any, ..self }
56    }
57
58    pub fn with_shape<S: Into<ShapeFactoid>>(self, shape: S) -> InferenceFact {
59        InferenceFact { shape: shape.into(), ..self }
60    }
61
62    pub fn format_dt_shape(&self) -> String {
63        if !self.shape.open && self.shape.dims.len() == 0 {
64            self.datum_type
65                .concretize()
66                .map(|dt| format!("{dt:?}"))
67                .unwrap_or_else(|| "?".to_string())
68        } else {
69            format!(
70                "{:?},{}",
71                self.shape,
72                self.datum_type
73                    .concretize()
74                    .map(|dt| format!("{dt:?}"))
75                    .unwrap_or_else(|| "?".to_string())
76            )
77        }
78    }
79
80    pub fn dt_shape_from_tensor(t: &Tensor) -> InferenceFact {
81        InferenceFact::dt_shape(t.datum_type(), t.shape())
82    }
83
84    pub fn without_value(self) -> InferenceFact {
85        InferenceFact { value: GenericFactoid::Any, ..self }
86    }
87
88    /// Parse a fact spec: comma-separated dims, optionally followed by the
89    /// element type, as in `1,3,S,f32`. Dims are TDim expressions resolved
90    /// against `symbols`, and `_` leaves a dim unknown. Without a trailing
91    /// element type the fact constrains the shape only; an empty spec
92    /// constrains nothing. The rank is always closed.
93    pub fn from_spec(symbols: &SymbolScope, spec: &str) -> TractResult<InferenceFact> {
94        if spec.is_empty() {
95            return Ok(InferenceFact::default());
96        }
97        let splits = spec.split(',').collect::<Vec<_>>();
98        let last = splits.last().unwrap();
99        let (datum_type, shape) = if let Ok(dt) = last.parse::<DatumType>() {
100            (Some(dt), &splits[0..splits.len() - 1])
101        } else {
102            (None, &*splits)
103        };
104        let shape = ShapeFactoid::closed(
105            shape
106                .iter()
107                .map(|&s| {
108                    Ok(if s == "_" {
109                        GenericFactoid::Any
110                    } else {
111                        GenericFactoid::Only(parse_tdim(symbols, s)?)
112                    })
113                })
114                .collect::<TractResult<TVec<DimFact>>>()?,
115        );
116        if let Some(dt) = datum_type {
117            Ok(InferenceFact::dt_shape(dt, shape))
118        } else {
119            Ok(InferenceFact::shape(shape))
120        }
121    }
122}
123
124impl Factoid for InferenceFact {
125    type Concrete = Arc<Tensor>;
126
127    /// Tries to transform the fact into a concrete value.
128    fn concretize(&self) -> Option<Self::Concrete> {
129        self.value.concretize()
130    }
131
132    /// Tries to unify the fact with another fact of the same type.
133    fn unify(&self, other: &Self) -> TractResult<Self> {
134        let tensor = InferenceFact {
135            datum_type: self.datum_type.unify(&other.datum_type)?,
136            shape: self.shape.unify(&other.shape)?,
137            value: self.value.unify(&other.value)?,
138        };
139
140        trace!("Unifying {self:?} with {other:?} into {tensor:?}.");
141
142        Ok(tensor)
143    }
144}
145
146impl fmt::Debug for InferenceFact {
147    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
148        if let Some(t) = self.value.concretize() {
149            write!(formatter, "{t:?}")
150        } else {
151            write!(formatter, "{}", self.format_dt_shape())
152        }
153    }
154}
155
156use crate::infer::factoid::Factoid;
157
158impl Fact for InferenceFact {
159    fn to_typed_fact(&self) -> TractResult<Cow<'_, TypedFact>> {
160        Ok(Cow::Owned(TypedFact::try_from(self)?))
161    }
162
163    fn matches(&self, t: &Tensor, _symbols: Option<&SymbolValues>) -> TractResult<bool> {
164        if let Some(dt) = self.datum_type() {
165            if t.datum_type() != dt {
166                return Ok(false);
167            }
168        }
169        if let Some(shape) = self.shape.concretize() {
170            if *ShapeFact::from(t.shape()) != *shape {
171                return Ok(false);
172            }
173        }
174        if let Some(value) = self.value.concretize() {
175            if &*value != t {
176                return Ok(false);
177            }
178        }
179        Ok(true)
180    }
181
182    fn compatible_with(&self, other: &dyn Fact) -> bool {
183        if let Some(other) = other.downcast_ref::<Self>() {
184            self.unify(other).is_ok()
185        } else {
186            false
187        }
188    }
189
190    fn datum_type(&self) -> Option<DatumType> {
191        self.datum_type.concretize()
192    }
193}
194
195impl TryFrom<&InferenceFact> for TypedFact {
196    type Error = TractError;
197    fn try_from(fact: &InferenceFact) -> TractResult<TypedFact> {
198        if let (Some(datum_type), Some(shape)) =
199            (fact.datum_type.concretize(), fact.shape.concretize())
200        {
201            let shape = ShapeFact::from_dims(shape);
202            let konst = fact.value.concretize();
203            let uniform = konst.as_ref().and_then(|k| k.as_uniform()).map(Arc::new);
204            Ok(TypedFact {
205                datum_type,
206                shape,
207                konst,
208                uniform,
209                exotic_fact: None,
210                uniform_tdim: None,
211                region_of_interest: None,
212            })
213        } else {
214            bail!("Can not make a TypedFact out of {:?}", fact)
215        }
216    }
217}
218
219impl<'a> From<&'a InferenceFact> for InferenceFact {
220    fn from(t: &'a InferenceFact) -> InferenceFact {
221        t.clone()
222    }
223}
224
225impl<'a> From<&'a TypedFact> for InferenceFact {
226    fn from(t: &'a TypedFact) -> InferenceFact {
227        let mut fact = InferenceFact::dt_shape(t.datum_type, t.shape.iter());
228        if let Some(k) = &t.konst {
229            fact.value = Arc::clone(k).into();
230        }
231        fact
232    }
233}
234
235impl From<TypedFact> for InferenceFact {
236    fn from(t: TypedFact) -> InferenceFact {
237        InferenceFact::from(&t)
238    }
239}
240
241impl<'a> TryFrom<&'a Arc<Tensor>> for InferenceFact {
242    type Error = TractError;
243    fn try_from(t: &'a Arc<Tensor>) -> TractResult<InferenceFact> {
244        Ok(InferenceFact::from(&TypedFact::try_from(Arc::clone(t))?))
245    }
246}
247
248impl TryFrom<Arc<Tensor>> for InferenceFact {
249    type Error = TractError;
250    fn try_from(t: Arc<Tensor>) -> TractResult<InferenceFact> {
251        Ok(InferenceFact::from(&TypedFact::try_from(t)?))
252    }
253}
254
255impl From<Tensor> for InferenceFact {
256    fn from(t: Tensor) -> InferenceFact {
257        let mut fact = InferenceFact::dt_shape(t.datum_type(), t.shape());
258        fact.value = t.into_arc_tensor().into();
259        fact
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use super::*;
266
267    #[test]
268    fn from_spec() {
269        let symbols = SymbolScope::default();
270        let s = symbols.sym("S");
271        assert_eq!(
272            InferenceFact::from_spec(&symbols, "1,S,_,f32").unwrap(),
273            InferenceFact::dt_shape(
274                DatumType::F32,
275                ShapeFactoid::closed(tvec!(
276                    GenericFactoid::Only(1.to_dim()),
277                    GenericFactoid::Only(s.to_dim()),
278                    GenericFactoid::Any
279                ))
280            )
281        );
282        assert_eq!(
283            InferenceFact::from_spec(&symbols, "2,3").unwrap(),
284            InferenceFact::shape(shapefactoid!(2, 3))
285        );
286        assert_eq!(
287            InferenceFact::from_spec(&symbols, "f32").unwrap(),
288            InferenceFact::dt_shape(DatumType::F32, ShapeFactoid::closed(tvec!()))
289        );
290        assert_eq!(InferenceFact::from_spec(&symbols, "").unwrap(), InferenceFact::default());
291    }
292}
tract_hir/infer/fact.rs

tract_hir/infer/
fact.rs
