tract_core/model/
fact.rs

1//! Partial and complete tensor types representations.
2use crate::internal::*;
3use downcast_rs::Downcast;
4use std::fmt;
5use tract_data::dyn_eq::DynEq;
6use tract_linalg::block_quant::{BlockQuantFact, BlockQuantStorage};
7
8#[derive(Clone, PartialEq, Eq, Hash)]
9pub struct ShapeFact {
10    dims: TVec<TDim>,
11    concrete: Option<TVec<usize>>,
12}
13
14impl ShapeFact {
15    #[inline]
16    pub fn rank(&self) -> usize {
17        self.dims.len()
18    }
19
20    fn compute_concrete(&mut self) {
21        assert!(self.dims.iter().all(|d| d.as_i64().map(|d| d >= 0).unwrap_or(true)));
22        self.concrete =
23            self.dims.iter().map(|d| d.as_i64().map(|d| d as usize)).collect::<Option<TVec<_>>>()
24    }
25
26    /// Shape of the tensor, unless it has symbolic dimensions.
27    #[inline]
28    pub fn as_concrete(&self) -> Option<&[usize]> {
29        self.concrete.as_deref()
30    }
31
32    /// Do we have a symbol-less value ?
33    #[inline]
34    pub fn is_concrete(&self) -> bool {
35        self.concrete.is_some()
36    }
37
38    /// Convert the shape to an array of extended dimensions.
39    #[inline]
40    pub fn to_tvec(&self) -> TVec<TDim> {
41        self.dims.clone()
42    }
43
44    /// Compute the volume of the tensor.
45    #[inline]
46    pub fn volume(&self) -> TDim {
47        self.dims.iter().product()
48    }
49
50    #[inline]
51    pub fn eval(&self, values: &SymbolValues) -> TractResult<Cow<'_, ShapeFact>> {
52        if self.is_concrete() {
53            Ok(Cow::Borrowed(self))
54        } else {
55            Ok(Cow::Owned(self.iter().map(|d| d.eval(values)).collect::<ShapeFact>()))
56        }
57    }
58
59    /// Substitute symbols by TDim expressions in every dim of the shape.
60    /// Concrete shapes pass through unchanged.
61    #[inline]
62    pub fn substitute(
63        &self,
64        subs: &std::collections::HashMap<Symbol, TDim>,
65    ) -> TractResult<Cow<'_, ShapeFact>> {
66        if self.is_concrete() {
67            Ok(Cow::Borrowed(self))
68        } else {
69            Ok(Cow::Owned(
70                self.iter().map(|d| d.substitute_all(subs)).collect::<TractResult<ShapeFact>>()?,
71            ))
72        }
73    }
74
75    #[inline]
76    pub fn eval_to_usize(&self, values: &SymbolValues) -> TractResult<Cow<'_, TVec<usize>>> {
77        if let Some(c) = &self.concrete {
78            Ok(Cow::Borrowed(c))
79        } else {
80            Ok(Cow::Owned(
81                self.iter()
82                    .map(|d| d.eval_to_i64(values).map(|d| d as usize))
83                    .collect::<TractResult<TVec<_>>>()?,
84            ))
85        }
86    }
87
88    #[inline]
89    pub fn eval_to_isize(&self, values: &SymbolValues) -> TractResult<Cow<'_, TVec<isize>>> {
90        if let Some(c) = &self.concrete {
91            #[allow(unknown_lints, clippy::missing_transmute_annotations)]
92            // TVec<usize> -> TVec<isize>
93            Ok(unsafe { std::mem::transmute(Cow::Borrowed(c)) })
94        } else {
95            Ok(Cow::Owned(
96                self.iter()
97                    .map(|d| d.eval_to_i64(values).map(|d| d as isize))
98                    .collect::<TractResult<TVec<_>>>()?,
99            ))
100        }
101    }
102
103    pub fn from_dims<D: ToDim, T: IntoIterator<Item = D>>(it: T) -> ShapeFact {
104        let mut dims =
105            ShapeFact { dims: it.into_iter().map(|d| d.to_dim()).collect(), concrete: None };
106        dims.compute_concrete();
107        dims
108    }
109
110    pub fn dims(&self) -> &[TDim] {
111        self.dims.as_slice()
112    }
113
114    pub fn set(&mut self, ix: usize, dim: TDim) {
115        self.dims[ix] = dim;
116        self.compute_concrete();
117    }
118
119    pub fn insert_axis(&mut self, axis: usize) -> TractResult<()> {
120        self.dims.insert(axis, 1.into());
121        if let Some(concrete) = &mut self.concrete {
122            concrete.insert(axis, 1);
123        }
124        Ok(())
125    }
126
127    pub fn remove_axis(&mut self, axis: usize) -> TractResult<()> {
128        self.dims.remove(axis);
129        if let Some(concrete) = &mut self.concrete {
130            concrete.remove(axis);
131        } else {
132            self.compute_concrete();
133        };
134        Ok(())
135    }
136
137    pub fn compatible_with(&self, _other: &ShapeFact) -> bool {
138        if self.rank() == _other.rank() {
139            self.dims
140                .iter()
141                .zip(_other.dims.iter())
142                .all(|(dim, other_dim)| dim.compatible_with(other_dim))
143        } else {
144            false
145        }
146    }
147
148    pub fn scalar() -> ShapeFact {
149        let void: &[usize] = &[];
150        Self::from(void)
151    }
152
153    pub fn consistent(&self) -> TractResult<()> {
154        ensure!(
155            self.concrete == self.dims.iter().map(|d| d.as_usize()).collect::<Option<TVec<_>>>()
156        );
157        Ok(())
158    }
159}
160
161impl std::ops::Deref for ShapeFact {
162    type Target = [TDim];
163    fn deref(&self) -> &[TDim] {
164        &self.dims
165    }
166}
167
168impl<D: ToDim, T: IntoIterator<Item = D>> From<T> for ShapeFact {
169    fn from(it: T) -> ShapeFact {
170        ShapeFact::from_dims(it)
171    }
172}
173
174/// Type information about a tensor: shape, and element type, in various state
175/// of determination.
176pub trait Fact:
177    std::fmt::Debug + Downcast + dyn_clone::DynClone + DynEq + Send + Sync + 'static
178{
179    fn to_typed_fact(&self) -> TractResult<Cow<'_, TypedFact>>;
180
181    fn matches(&self, t: &Tensor, symbols: Option<&SymbolValues>) -> TractResult<bool> {
182        self.to_typed_fact()?.matches(t, symbols)
183    }
184
185    /// Ensure that self is same type as another fact or a subtype
186    fn compatible_with(&self, _other: &dyn Fact) -> bool;
187
188    fn datum_type(&self) -> Option<DatumType>;
189}
190
191impl_downcast!(Fact);
192dyn_clone::clone_trait_object!(Fact);
193dyn_eq::eq_trait_object!(Fact);
194
195impl<D: ToDim> std::iter::FromIterator<D> for ShapeFact {
196    fn from_iter<T: IntoIterator<Item = D>>(iter: T) -> Self {
197        ShapeFact::from_dims(iter.into_iter().map(|d| d.to_dim()))
198    }
199}
200
201impl fmt::Debug for ShapeFact {
202    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
203        use tract_itertools::Itertools;
204        write!(fmt, "{}", self.iter().join(","))
205    }
206}
207
208impl AsRef<[TDim]> for ShapeFact {
209    fn as_ref(&self) -> &[TDim] {
210        &self.dims
211    }
212}
213
214/// Fully determined tensor information for TypedModel.
215#[derive(Clone, PartialEq, Eq, Hash)]
216pub struct TypedFact {
217    /// tensor element type
218    pub datum_type: DatumType,
219    /// tensor shape
220    pub shape: ShapeFact,
221    /// optional constant value
222    pub konst: Option<Arc<Tensor>>,
223    /// optional uniform value
224    pub uniform: Option<Arc<Tensor>>,
225    /// optional exotic fact
226    pub exotic_fact: Option<Box<dyn ExoticFact>>,
227    /// Symbolic per-element value as a TDim expression, possibly involving
228    /// coordinate symbols 🎯0,🎯1,… and/or model symbols.
229    /// `None` means "unknown / not tracked".
230    pub uniform_tdim: Option<TDim>,
231    /// Boolean TDim expression in coordinate symbols defining which positions
232    /// in the tensor are relevant to downstream consumers.
233    /// `None` means "all positions matter" (no demand annotation).
234    pub region_of_interest: Option<TDim>,
235}
236
237impl TypedFact {
238    pub fn scalar<T>() -> TypedFact
239    where
240        T: Datum,
241    {
242        Self::dt_scalar(T::datum_type())
243    }
244
245    pub fn shape<T, S>(shape: S) -> TypedFact
246    where
247        T: Datum,
248        S: Into<ShapeFact>,
249    {
250        Self::dt_shape(T::datum_type(), shape)
251    }
252
253    pub fn shape_and_dt_of(t: &Tensor) -> TypedFact {
254        debug_assert!(
255            t.is_plain(),
256            "shape_and_dt_of called on exotic tensor, exotic_fact will be lost"
257        );
258        TypedFact {
259            datum_type: t.datum_type(),
260            shape: ShapeFact::from_dims(t.shape().iter().map(TDim::from)),
261            uniform: None,
262            konst: None,
263            exotic_fact: None,
264            uniform_tdim: None,
265            region_of_interest: None,
266        }
267    }
268
269    pub fn mem_size(&self) -> TDim {
270        self.shape.volume() * self.datum_type.size_of()
271            + self.exotic_fact().iter().flat_map(|it| it.buffer_sizes()).sum::<TDim>()
272    }
273
274    pub fn dt_scalar(datum_type: DatumType) -> TypedFact {
275        TypedFact {
276            datum_type,
277            shape: ShapeFact::scalar(),
278            konst: None,
279            uniform: None,
280            exotic_fact: None,
281            uniform_tdim: None,
282            region_of_interest: None,
283        }
284    }
285
286    /// Parse a fact spec: the dims, then the element type, comma-separated, as
287    /// in `1,80,S,f32`. Dims are TDim expressions resolved against `symbols`,
288    /// so a spec with no dim at all (`f32`) is a scalar. Every dim must be
289    /// given: there is no wildcard, and no rank inference.
290    pub fn from_spec(symbols: &SymbolScope, spec: &str) -> TractResult<TypedFact> {
291        let mut parts = spec.split(',').map(|s| s.trim()).collect::<TVec<_>>();
292        let datum_type =
293            parts.pop().and_then(|last| last.parse::<DatumType>().ok()).with_context(|| {
294                format!("A fact spec ends with its element type, as in 1,80,f32; got {spec:?}")
295            })?;
296        let dims =
297            parts.iter().map(|dim| symbols.parse_tdim(dim)).collect::<TractResult<TVec<TDim>>>()?;
298        Ok(Self::dt_shape(datum_type, ShapeFact::from_dims(dims)))
299    }
300
301    pub fn dt_shape<S>(datum_type: DatumType, shape: S) -> TypedFact
302    where
303        S: Into<ShapeFact>,
304    {
305        TypedFact {
306            datum_type,
307            shape: shape.into(),
308            konst: None,
309            uniform: None,
310            exotic_fact: None,
311            uniform_tdim: None,
312            region_of_interest: None,
313        }
314    }
315
316    pub fn rank(&self) -> usize {
317        if cfg!(debug_assertions) {
318            self.consistent().unwrap();
319        }
320        self.shape.rank()
321    }
322
323    fn format_dt_shape_nocheck(&self) -> String {
324        if self.shape.rank() > 0 {
325            format!("{:?},{:?}", self.shape, self.datum_type)
326        } else {
327            format!("{:?}", self.datum_type)
328        }
329    }
330
331    pub fn format_dt_shape(&self) -> String {
332        if cfg!(debug_assertions) {
333            self.consistent().unwrap()
334        }
335        self.format_dt_shape_nocheck()
336    }
337
338    pub fn consistent(&self) -> TractResult<()> {
339        self.shape.consistent()?;
340        if let Some(k) = &self.konst {
341            if !self.matches(k.as_ref(), None)? {
342                bail!("fact says {}, constant is {:?}", self.format_dt_shape_nocheck(), k);
343            }
344            if let Some(bqf) = self.exotic_fact().and_then(|of| of.downcast_ref::<BlockQuantFact>())
345                && let Some(bqs) = k.storage_as::<BlockQuantStorage>()
346            {
347                let inner_bqf =
348                    BlockQuantFact::new(dyn_clone::clone_box(bqs.format()), k.shape().into());
349                ensure!(&inner_bqf == bqf, "BlockQuantStorage fact mismatch");
350            }
351        }
352        if let Some(u) = &self.uniform
353            && self.datum_type != u.datum_type()
354        {
355            bail!("fact as uniform value {:?}, but is of type {:?}", u, self.datum_type);
356        }
357        if let (Some(u), Some(k)) = (self.uniform.as_deref(), self.konst.as_deref()) {
358            if let Some(k) = k.as_uniform() {
359                if &k != u {
360                    bail!(
361                        "Uniform value and uniform constant mismatch: value:{u:?}, uniform:{k:?}",
362                    );
363                }
364            } else {
365                bail!("Fact said to be uniform ({:?}) and equal to {:?} which is not.", u, k);
366            }
367        }
368        Ok(())
369    }
370
371    pub fn without_value(&self) -> Self {
372        let mut new = self.clone();
373        new.konst = None;
374        new.uniform = None;
375        new.uniform_tdim = None;
376        new.region_of_interest = None;
377        new
378    }
379
380    pub fn with_exotic_fact<O: Into<Box<dyn ExoticFact>>>(mut self, exotic_fact: O) -> Self {
381        self.exotic_fact = Some(exotic_fact.into());
382        self
383    }
384
385    pub fn exotic_fact(&self) -> Option<&dyn ExoticFact> {
386        self.exotic_fact.as_deref()
387    }
388
389    #[inline]
390    pub fn is_exotic(&self) -> bool {
391        self.exotic_fact.is_some()
392    }
393
394    #[inline]
395    pub fn is_plain(&self) -> bool {
396        self.exotic_fact.is_none()
397    }
398}
399
400impl Fact for TypedFact {
401    fn to_typed_fact(&self) -> TractResult<Cow<'_, TypedFact>> {
402        if cfg!(debug_assertions) {
403            self.consistent()?
404        }
405        Ok(Cow::Borrowed(self))
406    }
407
408    fn matches(&self, t: &Tensor, symbols: Option<&SymbolValues>) -> TractResult<bool> {
409        if self.datum_type != t.datum_type() || self.shape.len() != t.rank() {
410            return Ok(false);
411        }
412        for i in 0..t.rank() {
413            if let Some(dim) = self.shape[i]
414                .eval(symbols.unwrap_or(&SymbolValues::default()))
415                .as_i64()
416                .map(|d| d as usize)
417                && dim != t.shape()[i]
418            {
419                return Ok(false);
420            }
421        }
422        Ok(true)
423    }
424
425    fn compatible_with(&self, other: &dyn Fact) -> bool {
426        if cfg!(debug_assertions) {
427            self.consistent().unwrap()
428        }
429        if let Some(other) = other.downcast_ref::<Self>() {
430            if cfg!(debug_assertions) {
431                other.consistent().unwrap()
432            }
433            self.datum_type == other.datum_type
434                && self.shape.compatible_with(&other.shape)
435                && self
436                    .exotic_fact()
437                    .zip(other.exotic_fact())
438                    .map(|(a, b)| a.compatible_with(b))
439                    .unwrap_or(true)
440        } else {
441            false
442        }
443    }
444
445    fn datum_type(&self) -> Option<DatumType> {
446        Some(self.datum_type)
447    }
448}
449
450impl TryFrom<Tensor> for TypedFact {
451    type Error = TractError;
452    fn try_from(t: Tensor) -> TractResult<TypedFact> {
453        TypedFact::try_from(t.into_arc_tensor())
454    }
455}
456
457impl TryFrom<Arc<Tensor>> for TypedFact {
458    type Error = TractError;
459    fn try_from(t: Arc<Tensor>) -> TractResult<TypedFact> {
460        let exotic_fact = t.exotic_fact()?;
461        let uniform_tdim = if t.datum_type() == TDim::datum_type() && t.len() == 1 {
462            t.try_as_plain_ram().ok().and_then(|d| d.as_slice::<TDim>().ok()).map(|s| s[0].clone())
463        } else if t.len() == 1
464            && t.try_as_plain_ram().is_ok()
465            && (t.datum_type().is_integer() || t.datum_type().is::<bool>())
466        {
467            t.cast_to_scalar::<i64>().ok().map(TDim::Val)
468        } else {
469            None
470        };
471        Ok(TypedFact {
472            datum_type: t.datum_type(),
473            shape: ShapeFact::from_dims(t.shape().iter().map(TDim::from)),
474            uniform: t.as_uniform().map(Arc::new),
475            exotic_fact,
476            konst: Some(t),
477            uniform_tdim,
478            region_of_interest: None,
479        })
480    }
481}
482
483impl From<&TypedFact> for TypedFact {
484    fn from(fact: &TypedFact) -> TypedFact {
485        fact.clone()
486    }
487}
488
489impl<'a> TryFrom<&'a Arc<Tensor>> for TypedFact {
490    type Error = TractError;
491    fn try_from(t: &'a Arc<Tensor>) -> TractResult<TypedFact> {
492        TypedFact::try_from(Arc::clone(t))
493    }
494}
495
496impl fmt::Debug for TypedFact {
497    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
498        write!(fmt, "{:?},{:?}", self.shape, self.datum_type)?;
499        if self.is_exotic() {
500            if let Some(of) = &self.exotic_fact {
501                write!(fmt, " 🔍 {of:?} ")?
502            } else {
503                write!(fmt, " 🔍 <no exotic fact> ")?
504            }
505        }
506        if let Some(k) = &self.konst {
507            write!(fmt, "🟰 {k:?}")?
508        }
509        if let Some(u) = &self.uniform {
510            write!(fmt, " ◻️{u:?}")?
511        }
512        if let Some(u) = &self.uniform_tdim {
513            write!(fmt, " 📐{u}")?
514        }
515        if let Some(r) = &self.region_of_interest {
516            write!(fmt, " 🬳 {r}")?
517        }
518        Ok(())
519    }
520}
521
522pub trait DatumExt {
523    fn scalar_fact() -> TypedFact;
524    fn fact<S>(shape: S) -> TypedFact
525    where
526        S: Into<ShapeFact>;
527}
528
529impl<T: Datum> DatumExt for T {
530    #[allow(clippy::needless_borrow)]
531    fn scalar_fact() -> TypedFact {
532        TypedFact::shape::<Self, &[usize]>(&[])
533    }
534
535    fn fact<S>(shape: S) -> TypedFact
536    where
537        S: Into<ShapeFact>,
538    {
539        TypedFact::shape::<Self, _>(shape)
540    }
541}
542
543pub trait DatumTypeExt {
544    fn scalar_fact(&self) -> TypedFact;
545    fn fact<S>(&self, shape: S) -> TypedFact
546    where
547        S: Into<ShapeFact>;
548}
549
550impl DatumTypeExt for DatumType {
551    #[allow(clippy::needless_borrow)]
552    fn scalar_fact(&self) -> TypedFact {
553        TypedFact::dt_shape::<&[usize]>(*self, &[])
554    }
555
556    fn fact<S>(&self, shape: S) -> TypedFact
557    where
558        S: Into<ShapeFact>,
559    {
560        TypedFact::dt_shape(*self, shape)
561    }
562}
563
564#[cfg(test)]
565mod from_spec_tests {
566    use super::*;
567
568    #[test]
569    fn scalar() {
570        let fact = TypedFact::from_spec(&SymbolScope::default(), "f32").unwrap();
571        assert_eq!(fact, f32::fact([0usize; 0]));
572    }
573
574    #[test]
575    fn concrete() {
576        let fact = TypedFact::from_spec(&SymbolScope::default(), "1,80,f16").unwrap();
577        assert_eq!(fact, f16::fact([1, 80]));
578    }
579
580    #[test]
581    fn symbolic() {
582        let symbols = SymbolScope::default();
583        let s = symbols.sym("S");
584        let fact = TypedFact::from_spec(&symbols, "1,80,2*S,f32").unwrap();
585        assert_eq!(fact, f32::fact(&[1.into(), 80.into(), s.to_dim() * 2]));
586    }
587
588    #[test]
589    fn missing_datum_type() {
590        assert!(TypedFact::from_spec(&SymbolScope::default(), "1,80").is_err());
591    }
592}
tract_core/model/fact.rs

tract_core/model/
fact.rs
