fix shapefact optimisation

sonos · Sep 5, 2023 · d1d95f5 · d1d95f5
1 parent 30530cc
commit d1d95f5
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/core/src/model/fact.rs b/core/src/model/fact.rs
@@ -105,7 +105,9 @@ impl ShapeFact {
         self.dims.remove(axis);
         if let Some(concrete) = &mut self.concrete {
             concrete.remove(axis);
-        }
+        } else {
+            self.compute_concrete();
+        };
         Ok(())
     }
 
@@ -124,6 +126,14 @@ impl ShapeFact {
         let void: &[usize] = &[];
         Self::from(void)
     }
+
+    pub fn consistent(&self) -> TractResult<()> {
+        ensure!(
+            self.concrete
+                == self.dims.iter().map(|d| d.to_usize()).collect::<TractResult<TVec<_>>>().ok()
+        );
+        Ok(())
+    }
 }
 
 impl std::ops::Deref for ShapeFact {
@@ -241,6 +251,7 @@ impl TypedFact {
     }
 
     pub fn consistent(&self) -> TractResult<()> {
+        self.shape.consistent()?;
         if let Some(k) = &self.konst {
             if !self.matches(k.as_ref(), None)? {
                 bail!("fact says {}, constant is {:?}", self.format_dt_shape_nocheck(), k);