Fix benchmark labels, remove invalid AveragePool dilations attribute, add pow NaN/Inf fallback tests, and harden Reshape dimension validation in

  propagate_shapes

Author: millioner

rust/jstprove_circuits/benches/autotuner_strategies.rs

@@ -23,11 +23,14 @@ fn clear_cache() {
     }
 }
 
+fn model_name() -> String {
+    std::env::var("MODEL").unwrap_or_else(|_| "lenet".to_string())
+}
+
 fn model_path() -> std::path::PathBuf {
-    let model_name = std::env::var("MODEL").unwrap_or_else(|_| "lenet".to_string());
     Path::new(env!("CARGO_MANIFEST_DIR"))
         .join("../jstprove_remainder/models")
-        .join(format!("{model_name}.onnx"))
+        .join(format!("{}.onnx", model_name()))
 }
 
 /// Benchmark cold compile: the autotuner cache is cleared before every iteration,
@@ -47,7 +50,7 @@ fn bench_cold_compile(c: &mut Criterion) {
     group.sample_size(10);
     group.measurement_time(Duration::from_secs(120));
 
-    group.bench_function("lenet", |b| {
+    group.bench_function(model_name(), |b| {
         b.iter_batched(
             || {
                 clear_cache();
@@ -103,7 +106,7 @@ fn bench_warm_compile(c: &mut Criterion) {
     group.sample_size(10);
     group.measurement_time(Duration::from_secs(120));
 
-    group.bench_function("lenet", |b| {
+    group.bench_function(model_name(), |b| {
         b.iter_batched(
             || {
                 OnnxContext::set_all(arch.clone(), params.clone(), Some(wandb.clone()));

rust/jstprove_circuits/benches/layer_bench.rs

@@ -758,10 +758,6 @@ fn make_averagepool_metadata() -> (CircuitParams, Architecture, WANDB) {
                 Value::String("pads".into()),
                 Value::Array(vec![Value::from(0i64); 4]),
             ),
-            (
-                Value::String("dilations".into()),
-                Value::Array(vec![Value::from(1i64), Value::from(1i64)]),
-            ),
         ])),
         opset_version_number: 17,
     };

rust/jstprove_circuits/src/circuit_functions/hints/pow.rs

@@ -174,4 +174,32 @@ mod tests {
             "(-2)^3: got {result}, expected {expected}"
         );
     }
+
+    #[test]
+    fn pow_zero_zero() {
+        // 0^0 → 0.0_f64.powf(0.0) = NaN → fallback to 1.0 (scale_u64 as i64)
+        let scale: u64 = 1 << 18;
+        let x_q = 0i64;
+        let exp_q = 0i64;
+        let result = run_hint(x_q, exp_q, scale);
+        assert_eq!(
+            result, scale as i64,
+            "0^0: expected 1.0 encoded as {}, got {result}",
+            scale
+        );
+    }
+
+    #[test]
+    fn pow_negative_fractional_exponent() {
+        // (-2)^0.5 → (-2.0_f64).powf(0.5) = NaN → fallback to 1.0 (scale_u64 as i64)
+        let scale: u64 = 1 << 18;
+        let x_q = -2 * scale as i64;
+        let exp_q = scale as i64 / 2; // 0.5 in fixed-point
+        let result = run_hint(x_q, exp_q, scale);
+        assert_eq!(
+            result, scale as i64,
+            "(-2)^0.5: expected 1.0 encoded as {}, got {result}",
+            scale
+        );
+    }
 }

rust/jstprove_onnx/src/quantizer.rs

@@ -329,7 +329,7 @@ fn compute_max_bound_inner(graph: &LayerGraph) -> Result<f64> {
     // default alpha for Log-layer bound propagation (ln(2^18) ≈ 12.47).
     let default_alpha_f64 = (DEFAULT_SCALE_BASE as f64).powi(DEFAULT_SCALE_EXPONENT as i32);
 
-    let shapes = propagate_shapes(graph);
+    let shapes = propagate_shapes(graph)?;
 
     for name in &graph.input_names {
         bounds.insert(name.clone(), 1.0);
@@ -413,7 +413,7 @@ fn broadcast_two(a: &[usize], b: &[usize]) -> Vec<usize> {
 ///
 /// This is used by `compute_bounds` so that reduction ops (e.g. ReduceSum)
 /// can account for the actual number of elements being accumulated.
-fn propagate_shapes(graph: &LayerGraph) -> HashMap<String, Vec<usize>> {
+fn propagate_shapes(graph: &LayerGraph) -> Result<HashMap<String, Vec<usize>>> {
     let mut shapes: HashMap<String, Vec<usize>> = HashMap::new();
 
     // Seed with model-level input shapes.
@@ -545,43 +545,68 @@ fn propagate_shapes(graph: &LayerGraph) -> HashMap<String, Vec<usize>> {
                         let in_shape: &[usize] = input_shape.as_deref().unwrap_or(&[]);
                         let input_total: usize = in_shape.iter().product();
 
-                        // First pass: resolve 0 → copy and positives; mark -1 as None.
-                        let mut dims: Vec<Option<usize>> = raw
-                            .iter()
-                            .enumerate()
-                            .map(|(i, &d)| {
-                                if d == 0 {
-                                    if allowzero {
-                                        Some(0)
-                                    } else {
-                                        Some(in_shape.get(i).copied().unwrap_or(0))
-                                    }
-                                } else if d > 0 {
-                                    Some(d as usize)
+                        // First pass: resolve 0 → copy and positives; -1 → infer; other negatives are invalid.
+                        let mut dims: Vec<Option<usize>> = Vec::with_capacity(raw.len());
+                        for (i, &d) in raw.iter().enumerate() {
+                            if d == 0 {
+                                dims.push(if allowzero {
+                                    Some(0)
                                 } else {
-                                    None // -1 only
-                                }
-                            })
-                            .collect();
+                                    Some(in_shape.get(i).copied().unwrap_or(0))
+                                });
+                            } else if d > 0 {
+                                dims.push(Some(d as usize));
+                            } else if d == -1 {
+                                dims.push(None); // infer sentinel
+                            } else {
+                                anyhow::bail!(
+                                    "Reshape layer '{}': invalid dimension {} at index {} (only -1 is a valid infer sentinel)",
+                                    layer.name, d, i
+                                );
+                            }
+                        }
 
                         // Infer the single -1 dimension when input total is known.
-                        if input_total > 0 {
-                            let n_unknown = dims.iter().filter(|d| d.is_none()).count();
-                            if n_unknown == 1 {
-                                let known: usize = dims.iter().filter_map(|&d| d).product();
-                                if known > 0 {
-                                    let inferred = input_total / known;
-                                    for d in &mut dims {
-                                        if d.is_none() {
-                                            *d = Some(inferred);
-                                            break;
-                                        }
-                                    }
+                        let n_unknown = dims.iter().filter(|d| d.is_none()).count();
+                        if n_unknown == 1 && input_total > 0 {
+                            let known: usize = dims.iter().filter_map(|&d| d).product();
+                            if known == 0 {
+                                anyhow::bail!(
+                                    "Reshape layer '{}': cannot infer -1 dimension when known product is 0",
+                                    layer.name
+                                );
+                            }
+                            if input_total % known != 0 {
+                                anyhow::bail!(
+                                    "Reshape layer '{}': input total {} is not divisible by known dims product {}",
+                                    layer.name, input_total, known
+                                );
+                            }
+                            let inferred = input_total / known;
+                            for d in &mut dims {
+                                if d.is_none() {
+                                    *d = Some(inferred);
+                                    break;
                                 }
                             }
+                        } else if n_unknown > 1 {
+                            anyhow::bail!(
+                                "Reshape layer '{}': more than one -1 dimension is not allowed",
+                                layer.name
+                            );
                         }
 
-                        dims.into_iter().map(|d| d.unwrap_or(0)).collect()
+                        // All None should be resolved by now; any remaining None is an error.
+                        dims.into_iter()
+                            .map(|d| {
+                                d.ok_or_else(|| {
+                                    anyhow::anyhow!(
+                                        "Reshape layer '{}': could not infer -1 dimension (input shape unknown)",
+                                        layer.name
+                                    )
+                                })
+                            })
+                            .collect::<Result<Vec<usize>>>()?
                     } else {
                         input_shape.unwrap_or_default()
                     }
@@ -1145,7 +1170,7 @@ fn propagate_shapes(graph: &LayerGraph) -> HashMap<String, Vec<usize>> {
         }
     }
 
-    shapes
+    Ok(shapes)
 }
 
 fn compute_bounds(graph: &LayerGraph, config: &ScaleConfig) -> Result<HashMap<String, usize>> {
@@ -1154,7 +1179,7 @@ fn compute_bounds(graph: &LayerGraph, config: &ScaleConfig) -> Result<HashMap<St
     let mut n_bits_config = HashMap::new();
 
     // Pre-compute tensor shapes for reduction-size estimation.
-    let shapes = propagate_shapes(graph);
+    let shapes = propagate_shapes(graph)?;
 
     for name in &graph.input_names {
         bounds.insert(name.clone(), 1.0);