Fix LayerNorm signed division/range constraints, improve LogUp out-of-bounds diagnostics, and add missing flatten_rmpv_to_f64 helper

Author: millioner

compiler/circuit-std-rs/src/logup.rs

@@ -469,7 +469,10 @@ pub fn query_count_hint<F: Field>(inputs: &[F], outputs: &mut [F]) -> Result<(),
     for input in inputs {
         let query_id = input.to_u256().as_usize();
         if query_id >= count.len() {
-            return Err(Error::InternalError("query_id out of bounds".into()));
+            return Err(Error::InternalError(format!(
+                "query_id out of bounds: query_id={query_id}, table_len={}",
+                count.len()
+            )));
         }
         count[query_id] += 1;
     }

rust/jstprove_circuits/src/circuit_functions/layers/layer_norm.rs

@@ -129,22 +129,40 @@ impl<C: Config, Builder: RootAPI<C>> LayerOp<C, Builder> for LayerNormLayer {
         let norm_offset = api.constant(CircuitField::<C>::from_u256(U256::from(
             1u64 << max_norm_bits,
         )));
+        // Signed-safe division support:
+        // norm_unscaled can be negative, while unconstrained_int_div/mod operate
+        // over field representatives. Bias by a large power-of-two multiple of
+        // `scale` so the dividend is non-negative in the intended integer domain,
+        // then subtract the corresponding quotient bias.
+        let norm_bias_bits = max_norm_bits + self.scale_exponent as usize + 2;
+        let norm_bias_u256 = U256::from(1u8) << norm_bias_bits;
+        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,
         )));
+        let norm_var_tol_native = lane_size as u64 * self.scaling * self.scaling;
         let norm_var_tolerance = api.constant(CircuitField::<C>::from_u256(U256::from(
-            lane_size as u64 * self.scaling,
+            norm_var_tol_native,
         )));
-        let norm_var_tol_bits = (2 * lane_size as u64 * self.scaling)
+        let norm_var_tol_bits = (2 * norm_var_tol_native)
             .next_power_of_two()
             .trailing_zeros() as usize;
 
         let mean_tolerance =
             api.constant(CircuitField::<C>::from_u256(U256::from(lane_size as u64)));
         let mean_tol_bits = (2 * lane_size + 1).next_power_of_two().trailing_zeros() as usize;
 
-        let per_elem_tolerance = api.constant(CircuitField::<C>::from_u256(U256::from(3u64)));
-        let per_elem_tol_bits: usize = 3;
+        // y_q comes from floating-point hint arithmetic while norm_q is reconstructed
+        // from quantized intermediates; allow bounded slack at alpha^2 scale.
+        let per_elem_tol_native = self.scaling * self.scaling;
+        let per_elem_tolerance = api.constant(CircuitField::<C>::from_u256(U256::from(
+            per_elem_tol_native,
+        )));
+        let per_elem_tol_bits = (2 * per_elem_tol_native)
+            .next_power_of_two()
+            .trailing_zeros() as usize;
 
         let outer_size: usize = shape[..axis].iter().product();
         let flat_input: Vec<Variable> = x_input
@@ -192,13 +210,16 @@ impl<C: Config, Builder: RootAPI<C>> LayerOp<C, Builder> for LayerNormLayer {
                 let dev = api.sub(flat_input[start + i], mean_q);
                 let norm_unscaled = api.mul(dev, inv_std_q);
 
-                let norm_q = api.unconstrained_int_div(norm_unscaled, scale_var);
-                let norm_rem = api.unconstrained_mod(norm_unscaled, scale_var);
-                let recon = api.mul(norm_q, scale_var);
+                let norm_unscaled_biased = api.add(norm_unscaled, norm_bias_var);
+                let norm_q_biased = api.unconstrained_int_div(norm_unscaled_biased, scale_var);
+                let norm_rem = api.unconstrained_mod(norm_unscaled_biased, scale_var);
+                let recon = api.mul(norm_q_biased, scale_var);
                 let recon = api.add(recon, norm_rem);
-                api.assert_is_equal(recon, norm_unscaled);
+                api.assert_is_equal(recon, norm_unscaled_biased);
                 logup_ctx.range_check::<C, Builder>(api, norm_rem, self.scale_exponent as usize)?;
 
+                let norm_q = api.sub(norm_q_biased, norm_q_bias_var);
+
                 let norm_shifted = api.add(norm_q, norm_offset);
                 logup_ctx.range_check::<C, Builder>(api, norm_shifted, max_norm_bits + 1)?;
 

rust/jstprove_circuits/src/runner/main_runner.rs

@@ -1790,6 +1790,58 @@ fn natural_key_cmp(a: &str, b: &str) -> std::cmp::Ordering {
         .then_with(|| a.cmp(b))
 }
 
+fn flatten_rmpv_to_f64(val: &Value, out: &mut Vec<f64>) -> Result<(), RunError> {
+    match val {
+        Value::Array(arr) => {
+            for item in arr {
+                flatten_rmpv_to_f64(item, out)?;
+            }
+            Ok(())
+        }
+        Value::Map(entries) => {
+            let mut pairs: Vec<_> = entries
+                .iter()
+                .filter_map(|(k, v)| k.as_str().map(|s| (s.to_string(), v)))
+                .collect();
+            pairs.sort_by(|(a, _), (b, _)| natural_key_cmp(a, b));
+            for (_, v) in pairs {
+                flatten_rmpv_to_f64(v, out)?;
+            }
+            Ok(())
+        }
+        Value::Integer(n) => {
+            let i = n
+                .as_i64()
+                .ok_or_else(|| RunError::Deserialize("integer value out of i64 range".into()))?;
+            out.push(i as f64);
+            Ok(())
+        }
+        Value::F64(f) => {
+            if f.is_finite() {
+                out.push(*f);
+                Ok(())
+            } else {
+                Err(RunError::Deserialize(
+                    "non-finite f64 value in input".into(),
+                ))
+            }
+        }
+        Value::F32(f) => {
+            if f.is_finite() {
+                out.push(f64::from(*f));
+                Ok(())
+            } else {
+                Err(RunError::Deserialize(
+                    "non-finite f32 value in input".into(),
+                ))
+            }
+        }
+        other => Err(RunError::Deserialize(format!(
+            "unsupported input value type for debug assignment: {other:?}"
+        ))),
+    }
+}
+
 #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
 fn flatten_value_to_i64(val: &Value) -> Vec<i64> {
     match val {