tract-harvest-probe: check run inputs against resolved facts; Pow exponent must not grow the base

crates/tract-harvest-probe/src/plan.rs

@@ -149,6 +149,13 @@ impl Plan<'_> {
                     node: (*name).into(),
                     what: "no such graph input".into(),
                 })?;
+            if !g.inputs.contains(&p) {
+                return Err(Refusal::Runtime {
+                    node: (*name).into(),
+                    what: "not a graph input".into(),
+                });
+            }
+            check_against_fact(name, v, &g.facts[p])?;
             vals[p] = Some(v.clone());
         }
         for s in &self.steps {
@@ -230,3 +237,34 @@ pub fn coverage(
     }
     out
 }
+
+/// A run-time input must satisfy the facts the plan was resolved on: its
+/// dtype, rank and every known dimension. Otherwise the WHEN guards that held
+/// at resolution say nothing about this run.
+fn check_against_fact(name: &str, v: &Value, f: &crate::registry::Fact) -> Result<(), Refusal> {
+    use crate::registry::Dim;
+    let refuse = |what: String| Refusal::Runtime {
+        node: name.into(),
+        what,
+    };
+    if let Some(dt) = f.dtype {
+        if dt != v.dtype() {
+            return Err(refuse(format!(
+                "dtype {:?}, resolved for {dt:?}",
+                v.dtype()
+            )));
+        }
+    }
+    if let Some(shape) = &f.shape {
+        let got = v.shape();
+        let fits = got.len() == shape.len()
+            && got.iter().zip(shape).all(|(g, d)| match d {
+                Dim::Known(k) => g == k,
+                Dim::Symbol(_) => true,
+            });
+        if !fits {
+            return Err(refuse(format!("shape {got:?}, resolved for {shape:?}")));
+        }
+    }
+    Ok(())
+}

crates/tract-harvest-probe/src/registry.rs

@@ -281,11 +281,52 @@ fn guard_cast_identity(v: &NodeView) -> K3 {
     }
 }
 
+/// `Pow(x, e)` admitted as `Mul(x, x)` only when every exponent element is
+/// exactly 2.0 AND broadcasting `e` against `x` leaves `x`'s shape unchanged
+/// (the polyfill ignores `e`, so it cannot grow the output).
 fn guard_pow2(v: &NodeView) -> K3 {
-    match v.consts.get(1).copied().flatten().and_then(Value::f32s) {
-        Some(e) if !e.is_empty() && e.iter().all(|x| x.to_bits() == 2.0f32.to_bits()) => K3::True,
-        Some(_) => K3::False,
-        None => K3::Unknown,
+    let Some(e) = v.consts.get(1).copied().flatten() else {
+        return K3::Unknown;
+    };
+    match e.f32s() {
+        Some(x) if !x.is_empty() && x.iter().all(|x| x.to_bits() == 2.0f32.to_bits()) => {}
+        _ => return K3::False,
+    }
+    let Some(base) = v
+        .args
+        .first()
+        .copied()
+        .flatten()
+        .and_then(|f| f.shape.as_ref())
+    else {
+        return K3::Unknown;
+    };
+    exponent_keeps_shape(base, &e.shape())
+}
+
+/// Does broadcasting a tensor of `exp` shape against `base` leave `base`'s
+/// shape unchanged?
+fn exponent_keeps_shape(base: &[Dim], exp: &[usize]) -> K3 {
+    if exp.len() > base.len() {
+        return if exp.iter().take(exp.len() - base.len()).all(|&d| d == 1) {
+            exponent_keeps_shape(base, &exp[exp.len() - base.len()..])
+        } else {
+            K3::False
+        };
+    }
+    let mut unknown = false;
+    for (b, &e) in base.iter().rev().zip(exp.iter().rev()) {
+        match b {
+            _ if e == 1 => {}
+            Dim::Known(k) if *k == e => {}
+            Dim::Known(_) => return K3::False,
+            Dim::Symbol(_) => unknown = true,
+        }
+    }
+    if unknown {
+        K3::Unknown
+    } else {
+        K3::True
     }
 }
 
@@ -774,3 +815,73 @@ pub fn registry() -> Registry {
         ],
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn pow_view<'a>(node: &'a SourceNode, base: &'a Fact, exp: &'a Value) -> NodeView<'a> {
+        NodeView {
+            node,
+            args: vec![Some(base), None],
+            consts: vec![None, Some(exp)],
+        }
+    }
+
+    fn twos(shape: &[usize]) -> Value {
+        Value::F32(ArrayD::from_elem(IxDyn(shape), 2.0))
+    }
+
+    #[test]
+    fn pow2_requires_the_exponent_not_to_grow_the_base() {
+        let node = SourceNode {
+            name: "p".into(),
+            op_type: "Pow".into(),
+            domain: String::new(),
+            inputs: vec!["x".into(), "e".into()],
+            outputs: vec!["y".into()],
+            attributes: vec![],
+            attrs: vec![],
+        };
+        let base = |s: &[usize]| Fact {
+            dtype: Some(DType::F32),
+            shape: Some(s.iter().map(|&d| Dim::Known(d)).collect()),
+        };
+        // scalar and broadcast-neutral exponents
+        for (b, e) in [
+            (&[1, 10, 1024][..], &[][..]),
+            (&[1, 10, 1024], &[1]),
+            (&[2, 3], &[1, 3]),
+            (&[2, 3], &[1, 1, 1]),
+        ] {
+            assert_eq!(
+                guard_pow2(&pow_view(&node, &base(b), &twos(e))),
+                K3::True,
+                "{b:?} ^ {e:?}"
+            );
+        }
+        // the exponent would grow the output: the polyfill must not run
+        for (b, e) in [(&[1][..], &[2][..]), (&[2, 1], &[2, 3]), (&[3], &[2, 3])] {
+            assert_eq!(
+                guard_pow2(&pow_view(&node, &base(b), &twos(e))),
+                K3::False,
+                "{b:?} ^ {e:?}"
+            );
+        }
+        // a symbolic base dim against a non-1 exponent dim is UNKNOWN
+        let sym = Fact {
+            dtype: Some(DType::F32),
+            shape: Some(vec![Dim::Symbol("S".into())]),
+        };
+        assert_eq!(guard_pow2(&pow_view(&node, &sym, &twos(&[4]))), K3::Unknown);
+        // a non-2 exponent is FALSE whatever its shape
+        assert_eq!(
+            guard_pow2(&pow_view(
+                &node,
+                &base(&[3]),
+                &Value::F32(ArrayD::from_elem(IxDyn(&[]), 3.0))
+            )),
+            K3::False
+        );
+    }
+}

crates/tract-harvest-probe/tests/tiny.rs

@@ -385,3 +385,35 @@ fn unknown_shapes_refuse() {
         other => panic!("{:?}", other.map(|_| ())),
     }
 }
+
+/// An input outside the facts the plan was resolved on is refused before any
+/// node runs: `X` resolved as `[7]` (rank 1, what `ReduceSum`'s WHEN admitted)
+/// supplied as `[2,3]`, or with the wrong dtype.
+#[test]
+fn inputs_outside_the_resolved_facts_are_refused() {
+    let reg = registry();
+    let (fx, p) = plan("base", &reg);
+    let p = p.unwrap();
+    let mut ins = inputs(None);
+    ins[0].1 = Value::I32(ArrayD::from_shape_vec(IxDyn(&[2, 3]), vec![1; 6]).unwrap());
+    match p.run(&fx.store, &ins) {
+        Err(Refusal::Runtime { node, what }) => {
+            assert_eq!(node, "X");
+            assert!(what.contains("shape [2, 3]"), "{what}");
+        }
+        other => panic!("{:?}", other.map(|_| ())),
+    }
+    let mut ins = inputs(None);
+    ins[1].1 = Value::I32(ArrayD::from_shape_vec(IxDyn(&[2, 3]), vec![1; 6]).unwrap());
+    assert!(matches!(
+        p.run(&fx.store, &ins),
+        Err(Refusal::Runtime { .. })
+    ));
+    // an initializer is not an input the caller may overwrite
+    let mut ins = inputs(None);
+    ins.push(("B", f32_value(&[3], &[0; 3])));
+    assert!(matches!(
+        p.run(&fx.store, &ins),
+        Err(Refusal::Runtime { .. })
+    ));
+}