Fail fast on missing MODEL path, pin logup_chunk_bits in pipeline benches, add ConvTranspose zero-stride and bias ndim guards, and fix LayerNorm variance

millioner · millioner · commit d908f2fb040f · 2026-04-08T22:25:11.000+02:00
tolerance overflow and zero-bits crash
diff --git a/rust/jstprove_circuits/benches/autotuner_strategies.rs b/rust/jstprove_circuits/benches/autotuner_strategies.rs
@@ -37,9 +37,12 @@ fn model_path() -> std::path::PathBuf {
 /// forcing a full sweep each time.
 fn bench_cold_compile(c: &mut Criterion) {
     let path = model_path();
-    if !path.exists() {
-        return;
-    }
+    assert!(
+        path.exists(),
+        "MODEL='{}' resolved to '{}' — ONNX file not found; set the MODEL env var to a model in jstprove_remainder/models/",
+        model_name(),
+        path.display()
+    );
 
     let metadata = expander_metadata::generate_from_onnx(&path).unwrap();
     // Force logup_chunk_bits = None so compile_bn254 always invokes the autotuner
@@ -80,9 +83,12 @@ fn bench_cold_compile(c: &mut Criterion) {
 /// every subsequent iteration hits the cache and skips the sweep.
 fn bench_warm_compile(c: &mut Criterion) {
     let path = model_path();
-    if !path.exists() {
-        return;
-    }
+    assert!(
+        path.exists(),
+        "MODEL='{}' resolved to '{}' — ONNX file not found; set the MODEL env var to a model in jstprove_remainder/models/",
+        model_name(),
+        path.display()
+    );
 
     // Pre-populate the cache with one cold compile.
     {
diff --git a/rust/jstprove_circuits/benches/pipeline_ecc.rs b/rust/jstprove_circuits/benches/pipeline_ecc.rs
@@ -25,12 +25,18 @@ fn model_path() -> std::path::PathBuf {
 
 fn bench_bn254(c: &mut Criterion) {
     let path = model_path();
-    if !path.exists() {
-        return;
-    }
+    assert!(
+        path.exists(),
+        "MODEL='{}' resolved to '{}' — ONNX file not found; set the MODEL env var to a model in jstprove_remainder/models/",
+        model_name(),
+        path.display()
+    );
 
     let metadata = expander_metadata::generate_from_onnx(&path).unwrap();
-    let params = metadata.circuit_params.clone();
+    // Pin a fixed chunk width so every compile iteration runs in the same warm-cache
+    // regime without triggering the adaptive autotuner sweep on the first iteration.
+    let mut params = metadata.circuit_params.clone();
+    params.logup_chunk_bits = params.logup_chunk_bits.or(Some(12));
     let arch = metadata.architecture.clone();
     let wandb = metadata.wandb.clone();
     OnnxContext::set_all(arch.clone(), params.clone(), Some(wandb.clone()));
@@ -104,13 +110,17 @@ fn bench_bn254(c: &mut Criterion) {
 
 fn bench_goldilocks(c: &mut Criterion) {
     let path = model_path();
-    if !path.exists() {
-        return;
-    }
+    assert!(
+        path.exists(),
+        "MODEL='{}' resolved to '{}' — ONNX file not found; set the MODEL env var to a model in jstprove_remainder/models/",
+        model_name(),
+        path.display()
+    );
 
     let metadata =
         expander_metadata::generate_from_onnx_for_field(&path, N_BITS_GOLDILOCKS, None).unwrap();
-    let params = metadata.circuit_params.clone();
+    let mut params = metadata.circuit_params.clone();
+    params.logup_chunk_bits = params.logup_chunk_bits.or(Some(12));
     let arch = metadata.architecture.clone();
     let wandb = metadata.wandb.clone();
     OnnxContext::set_all(arch.clone(), params.clone(), Some(wandb.clone()));
@@ -184,13 +194,17 @@ fn bench_goldilocks(c: &mut Criterion) {
 
 fn bench_goldilocks_basefold(c: &mut Criterion) {
     let path = model_path();
-    if !path.exists() {
-        return;
-    }
+    assert!(
+        path.exists(),
+        "MODEL='{}' resolved to '{}' — ONNX file not found; set the MODEL env var to a model in jstprove_remainder/models/",
+        model_name(),
+        path.display()
+    );
 
     let metadata =
         expander_metadata::generate_from_onnx_for_field(&path, N_BITS_GOLDILOCKS, None).unwrap();
-    let params = metadata.circuit_params.clone();
+    let mut params = metadata.circuit_params.clone();
+    params.logup_chunk_bits = params.logup_chunk_bits.or(Some(12));
     let arch = metadata.architecture.clone();
     let wandb = metadata.wandb.clone();
     OnnxContext::set_all(arch.clone(), params.clone(), Some(wandb.clone()));
diff --git a/rust/jstprove_circuits/src/circuit_functions/layers/conv_transpose.rs b/rust/jstprove_circuits/src/circuit_functions/layers/conv_transpose.rs
@@ -181,6 +181,13 @@ impl<C: Config, Builder: RootAPI<C>> LayerOp<C, Builder> for ConvTransposeLayer
         let mut res = if bias.is_empty() {
             ArrayD::from_elem(vec![n_batch, c_out, out_h, out_w], zero)
         } else {
+            if bias.ndim() != 1 {
+                return Err(LayerError::InvalidShape {
+                    layer: LayerKind::ConvTranspose,
+                    msg: format!("bias must be 1-D, got {}D", bias.ndim()),
+                }
+                .into());
+            }
             if bias.shape()[0] != c_out {
                 return Err(LayerError::InvalidShape {
                     layer: LayerKind::ConvTranspose,
@@ -363,6 +370,18 @@ impl<C: Config, Builder: RootAPI<C>> LayerOp<C, Builder> for ConvTransposeLayer
             }
         }
 
+        // Guard against zero strides: the apply method uses stride as a divisor.
+        if strides.contains(&0) {
+            return Err(LayerError::UnsupportedConfig {
+                layer: LayerKind::ConvTranspose,
+                msg: format!(
+                    "layer '{}': strides must all be > 0, got {:?}",
+                    layer.name, strides
+                ),
+            }
+            .into());
+        }
+
         let conv_transpose = Self {
             weights,
             bias,
diff --git a/rust/jstprove_circuits/src/circuit_functions/layers/layer_norm.rs b/rust/jstprove_circuits/src/circuit_functions/layers/layer_norm.rs
@@ -139,20 +139,25 @@ impl<C: Config, Builder: RootAPI<C>> LayerOp<C, Builder> for LayerNormLayer {
         let norm_bias_var = api.constant(CircuitField::<C>::from_u256(norm_bias_u256));
         let norm_q_bias_u256 = U256::from(1u8) << (norm_bias_bits - self.scale_exponent as usize);
         let norm_q_bias_var = api.constant(CircuitField::<C>::from_u256(norm_q_bias_u256));
-        let n_alpha_sq = api.constant(CircuitField::<C>::from_u256(U256::from(
-            lane_size as u64 * self.scaling * self.scaling,
-        )));
+        // Compute N·α² in U256 to avoid u64 overflow when lane_size is large or
+        // scale_exponent approaches the u64-squaring limit (guard in build caps at < 32,
+        // but lane_size * 2^(2*exp) can still exceed u64 with exp=31 and lane_size≥4).
+        let n_alpha_sq_u256 =
+            U256::from(lane_size as u64) * U256::from(self.scaling) * U256::from(self.scaling);
+        let n_alpha_sq = api.constant(CircuitField::<C>::from_u256(n_alpha_sq_u256));
         // Use half of n_alpha_sq as the variance tolerance so that the constraint
         // norm_sq_sum ∈ [N·α²/2, 3N·α²/2) is non-trivially enforced.  With the
         // full N·α² tolerance the range-check reduced to just checking norm_sq_sum
         // is non-negative, which a cheating prover could trivially satisfy.
-        let norm_var_tol_native = (lane_size as u64 * self.scaling * self.scaling) / 2;
-        let norm_var_tolerance = api.constant(CircuitField::<C>::from_u256(U256::from(
-            norm_var_tol_native,
-        )));
-        let norm_var_tol_bits = (2 * norm_var_tol_native)
-            .next_power_of_two()
-            .trailing_zeros() as usize;
+        let norm_var_tol_u256 = n_alpha_sq_u256 >> 1u32; // / 2
+        let norm_var_tolerance = api.constant(CircuitField::<C>::from_u256(norm_var_tol_u256));
+        // ceil(log2(2 * norm_var_tol)) using U256; 0 is handled separately below.
+        let norm_var_tol_bits: usize = if norm_var_tol_u256 == U256::from(0u64) {
+            0
+        } else {
+            let two_tol = norm_var_tol_u256 * U256::from(2u64);
+            (256u32 - (two_tol - U256::from(1u64)).leading_zeros()) as usize
+        };
 
         let mean_tolerance =
             api.constant(CircuitField::<C>::from_u256(U256::from(lane_size as u64)));
@@ -251,8 +256,14 @@ impl<C: Config, Builder: RootAPI<C>> LayerOp<C, Builder> for LayerNormLayer {
             }
 
             let var_diff = api.sub(norm_sq_sum, n_alpha_sq);
-            let var_shifted = api.add(var_diff, norm_var_tolerance);
-            logup_ctx.range_check::<C, Builder>(api, var_shifted, norm_var_tol_bits)?;
+            if norm_var_tol_bits == 0 {
+                // Tolerance is 0 (degenerate: scale_exponent = 0 and lane_size = 1);
+                // range_check(_, 0) is an error, so assert equality directly.
+                api.assert_is_equal(var_diff, zero_var);
+            } else {
+                let var_shifted = api.add(var_diff, norm_var_tolerance);
+                logup_ctx.range_check::<C, Builder>(api, var_shifted, norm_var_tol_bits)?;
+            }
 
             flat_output.extend_from_slice(y_vars);
         }
