From a093d65d2ebe6a69c495d768c08509989881bb41 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 07:30:49 +0000 Subject: [PATCH] tract-harvest-probe: check run inputs against resolved facts; Pow exponent must not grow the base Codex P2 (#1850): run() bound caller inputs without checking the facts the WHEN guards held on, so X as [2,3] reached the rank-1 ReduceSum DO and was summed flat. Inputs are now refused unless they are graph inputs with the resolved dtype, rank and known dims. Codex P2: guard_pow2 checked only the exponent's values; do_pow2 ignores the exponent, so a broadcasting exponent ([1] ^ [2]) gave the wrong output shape. WHEN now also requires broadcasting e against x to leave x's shape unchanged (FALSE, or UNKNOWN for a symbolic base dim). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XAqKMMjKnig6MXTd41Lf2v --- crates/tract-harvest-probe/src/plan.rs | 38 +++++++ crates/tract-harvest-probe/src/registry.rs | 119 ++++++++++++++++++++- crates/tract-harvest-probe/tests/tiny.rs | 32 ++++++ 3 files changed, 185 insertions(+), 4 deletions(-) diff --git a/crates/tract-harvest-probe/src/plan.rs b/crates/tract-harvest-probe/src/plan.rs index 9bfe2839d..43b9dd1ea 100644 --- a/crates/tract-harvest-probe/src/plan.rs +++ b/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(()) +} diff --git a/crates/tract-harvest-probe/src/registry.rs b/crates/tract-harvest-probe/src/registry.rs index 3e34c9ff0..5be499690 100644 --- a/crates/tract-harvest-probe/src/registry.rs +++ b/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 + ); + } +} diff --git a/crates/tract-harvest-probe/tests/tiny.rs b/crates/tract-harvest-probe/tests/tiny.rs index 92ec53906..88df9721a 100644 --- a/crates/tract-harvest-probe/tests/tiny.rs +++ b/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", ®); + 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 { .. }) + )); +}