Refactor circuit functions into gadgets/hints, improve quantization docs, and fix lint

[tmfreiberg]

Description

This PR refactors the JSTprove circuits crate into clearer submodules (gadgets, hints, layers, utils), adds detailed documentation for core circuit gadgets (particularly max/min/clip and range checks), and brings both the Rust and Python code in line with our linting/pre-commit configuration. Functionally, behavior should be unchanged; the focus is structure, clarity, and tooling hygiene.

---

Rust changes (jstprove_circuits)

1. Circuit function reorganisation

• Reorganised rust/jstprove_circuits/src/circuit_functions into:

  • gadgets/ – reusable constrained circuit gadgets (e.g., max/min/clip, range checks).
  • hints/ – unconstrained witness helpers and LogUp-related hints.
  • layers/ – ONNX-style layer frontends (e.g., Clip, MaxPool) that call into gadgets.
  • utils/ – shared utilities (e.g., quantization/rescaling helpers).

This is largely a file/move and naming refactor; the intent is to align the directory structure with the mental model we use in design docs (unconstrained hints → gadgets → layers).

2. Max/Min/Clip gadgets + range checks

• Implemented / documented constrained max/min/clip gadgets that:

  • Use LogUp-based range proofs rather than ad-hoc MSB extraction.
  • Enforce "selected value is one of the inputs" via product-of-differences constraints.

• Added detailed doc comments on:

  • The semantics (how MaxPool, MaxLayer, and MinLayer use these gadgets).
  • The LogUp-based argument (shift, range checks, product constraint, and field assumptions).

• Added a LogUp range-check context:

  • LogupRangeCheckContext and helper functions to amortise base-2^chunk_bits range checks across many values.
  • Wrapper helper for the common [0, 2^{n_bits} - 1] unsigned range-check case.

These gadgets are intended to be the canonical primitives used by max/min/clip layers and quantization code.

3. Unconstrained helpers (hints)

• Added / documented unconstrained helpers in hints/, including:

  • unconstrained_to_bits for extracting LSBs via unconstrained bitwise ops.

  • unconstrained_max, unconstrained_min, and unconstrained_clip as witness-level helpers that:

    • Use unconstrained comparisons/selection.
    • Do not enforce any constraints themselves.

• These are explicit "witness only" APIs; all correctness is expected to be enforced by constrained gadgets.

4. Quantization utilities (utils/quantization.rs)

• Documented the fixed-point rescaling protocol in detail:

  • Assumptions: a ≈ α·x, b ≈ α·y with α = 2^κ.
  • Computation of q = floor((c + α·S)/α) − S with S = 2^s.
  • Use of q^♯ = q + S and remainder r with LogUp range checks on both.
  • Optional ReLU(q) = max(q, 0) implemented via constrained_max.

• Introduced RescalingContext to:

  • Precompute integer powers-of-two (α, S, α·S) safely (u32/U256).
  • Lift these constants once into circuit Variables for reuse.

• rescale:

  • Uses RescalingContext + shared LogupRangeCheckContext.
  • Enforces shifted_dividend = α·q^♯ + r and applies LogUp range checks on r and q^♯.
  • Optionally applies ReLU using the same shift exponent as the rescaling context.

• rescale_array:

  • Applies rescale elementwise over an ArrayD<Variable>, using a single shared LogUp context for all range checks.
  • Converts between ArrayD and Vec with explicit error handling for shape mismatches.

These changes are intended to make the quantization / rescaling logic auditable and reusable, without changing semantics.

5. Clippy + doc cleanups

• Addressed some Clippy warnings in the Rust crate, including:

  • slienced Clippy at lib.rs level on nits related to doc_markdown, doc_overindented_list_items, doc_comment_double_space_linebreaks.
  • Adding # Errors / # Panics sections where appropriate.
  • Adding #[must_use] on context constructors that are intended to be used.

• In a few places, explicit #[allow(...)] or # noqa-style annotations were added where the lint conflicts with existing API or semantics (rather than forcing an invasive change).

---

Python & tooling changes

1. Pre-commit and lint alignment

• Ran pre-commit across the repo and fixed:

  • End-of-file newlines and trailing whitespace.
  • black formatting on a few Python modules.
  • ruff nits in selected files.

2. Python code cleanups (high level)

Some representative changes (non-exhaustive; see individual diffs):

• python/core/model_processing/onnx_custom_ops/__init__.py

  • Use pathlib.Path instead of os.path.dirname.
  • Switch __all__.append(...) to __all__ += [...] for better type-checker compatibility.
  • Add type hints for __all.

• Quantizer & errors modules:

  • Minor signature/type-hint fixes.
  • Localised ruff suppressions for rules that would be API-breaking or not worth "perfect" compliance.

• Tests and helpers:

  • Removed unused local variables and commented-out debug prints.
  • Broke a few long lines to satisfy E501.
  • Updated one open("model.onnx", "wb") call to Path.open().

• Template / scratch files:

  • Some template/scratch modules still exist but are now consistent with current linting, or explicitly annotated where appropriate.

No Python behavior related to core circuit functionality is intended to change; these are formatting and hygiene changes.

---

Testing

I ran (and everything passed):

• ✅ cargo build --release
• ✅ Python test suite (e.g., pytest --model lenet etc. across the repo).
• ✅ jst compile, witness, prove, verify works for lenet

Related Issue

Type of Change

• Bug fix (non-breaking)
• New feature (non-breaking)
• Breaking change (fix/feature causing existing functionality to break)
• Refactor (non-functional changes)
• Documentation update

Checklist

• Code follows project patterns
• Tests added/updated (if applicable)
• Documentation updated (if applicable)
• Self-review of code
• All tests pass locally
• Linter passes locally

Deployment Notes

Additional Comments

• The Rust changes are mostly organizational and documentation changes with some light refactoring into contexts and helper functions. There should be no observable behavioral change in existing circuits.

• The most "substantive" pieces worth a closer look are:

  • The LogUp-based max/min/clip gadget implementations and their doc comments.
  • The RescalingContext / rescale protocol and error handling.

• Most Python changes are automated formatting or straightforward lint-driven edits; they're included to keep CI and pre-commit green going forward.

Summary by CodeRabbit

• New Features
  • ONNX quantization for Clip, Max, and Min with new constraint gadgets and a hint registry for stronger verification and witness generation.
• Bug Fixes
  • Fixed toolchain TOML syntax and minor test/data formatting issues.
• Code Quality
  • Modernized type hints, refined public APIs, and added public modules for gadgets, hints, IO, and runner.
• Tests
  • Added providers/specs for Clip/Max/Min and standardized test markers.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Actionable comments posted: 11

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

rust/jstprove_circuits/src/circuit_functions/layers/maxpool.rs (1)

449-463: The bounds checking logic is correct; remove unreachable dead code from the guard conditions.

The pattern Ok(val) if h < 0 || val < height => val correctly uses the guard to validate that val < height (in-bounds), returning the index when valid and continuing (skipping) when val >= height. Since usize::try_from(h) only succeeds when h >= 0, the h < 0 check is unreachable in the Ok branch and should be removed for clarity.

The same applies to the w check on line 457.

                     for h in (hstart..hend).step_by(dilation_h) {
-                        let h_usize = match usize::try_from(h) {
-                            Ok(val) if h < 0 || val < height => val,
+                        let h_usize = match usize::try_from(h) {
+                            Ok(val) if val < height => val,
                             _ => continue,
                         };

                         for w in (wstart..wend).step_by(dilation_w) {
-                            let w_usize = match usize::try_from(w) {
-                                Ok(val) if w < 0 || val < width => val,
+                            let w_usize = match usize::try_from(w) {
+                                Ok(val) if val < width => val,
                                 _ => continue,
                             };

rust/jstprove_circuits/src/circuit_functions/utils/quantization.rs (1)

139-144: Copy-paste error in error handling: uses scaling_exponent instead of shift_exponent.

The error message for shift_ overflow incorrectly references scaling_exponent instead of shift_exponent on line 141.

         let shift_ = 1u32.checked_shl(shift_exponent_u32).ok_or(
             RescaleError::ShiftExponentTooLargeError {
-                exp: scaling_exponent,
+                exp: shift_exponent,
                 type_name: "u32",
             },
         )?;

🧹 Nitpick comments (18)

python/core/model_processing/onnx_custom_ops/__init__.py (2)

13-13: Consider explicit string conversion for clarity.

While Path objects work with pkgutil.iter_modules() via the __fspath__() protocol, explicit conversion makes the intent clearer and ensures compatibility.

Apply this diff:

-for _, module_name, is_pkg in pkgutil.iter_modules([package_dir]):
+for _, module_name, is_pkg in pkgutil.iter_modules([str(package_dir)]):

---

16-16: Remove redundant string cast and unused noqa directive.

The str() cast is redundant since module_name from pkgutil.iter_modules() is already a string. The noqa: PYI056 directive is unused because PYI056 applies only to stub files (.pyi), not regular Python files.

Apply this diff:

-        __all__.append(str(module_name))  # noqa: PYI056
+        __all__.append(module_name)

python/core/model_templates/circuit_template.py (1)

53-67: Avoid duplicating input dict construction in get_outputs

Logic is correct (output = input_a + input_b), but the local reconstruction of inputs duplicates get_inputs. To keep things DRY and robust if required_keys / defaults evolve, consider:

-        if inputs is None:
-            inputs = {
-                "input_a": self.input_a,
-                "input_b": self.input_b,
-                "nonce": self.nonce,
-            }
+        if inputs is None:
+            inputs = self.get_inputs()

This keeps a single source of truth for the input structure.

python/tests/circuit_e2e_tests/circuit_model_developer_test.py (1)

42-42: Optional: Marker formatting is cleaner without parentheses.

The change from @pytest.mark.e2e() to @pytest.mark.e2e is a stylistic improvement. Both forms are valid, but the latter is more concise when no arguments are passed.

Also applies to: 50-50, 96-96, 152-152, 246-246, 346-346, 428-428, 474-474, 562-562, 673-673, 797-797, 944-944, 1070-1070

python/core/model_processing/onnx_quantizer/exceptions.py (1)

154-154: Remove unused noqa directive.

The # noqa: FBT001 directive is unused because the FBT001 rule (boolean positional arguments) is not enabled in your linting configuration.

Apply this diff:

-        value: str | float | bool | None,  # noqa: FBT001
+        value: str | float | bool | None,

python/tests/onnx_quantizer_tests/layers/min_config.py (1)

40-85: Comprehensive test coverage; consider improving type hint.

The test specifications provide excellent coverage including basic operations, broadcasting, initializers, edge cases, and e2e tests. However, consider making the return type more specific:

-    def get_test_specs(self) -> list:
+    def get_test_specs(self) -> list[LayerTestSpec]:

This improves type safety and matches the pattern used in other config providers.

rust/jstprove_circuits/src/runner/main_runner.rs (1)

291-302: Reuse the same LogUp hint registry in debug_witness

debug_witness currently builds the LogUp hint registry twice even though the comment says “Build LogUp registry once”:

• logup_hints is built and cloned into debug_eval.
• A second hint_registry is built for solve_witness_with_hints.

Because debug_eval receives a clone by value, you can safely reuse logup_hints for the solver and avoid the extra construction:

-    // Build LogUp registry once
-    let logup_hints = build_logup_hint_registry::<CircuitField<C>>();
-
-    // Use it for the frontend debug evaluation
-    debug_eval(&circuit, &assignment, logup_hints.clone());
-
-    // And for the witness solver
-    let hint_registry = build_logup_hint_registry::<CircuitField<C>>();
-
-    let witness = witness_solver
-        .solve_witness_with_hints(&assignment, &hint_registry)
+    // Build LogUp registry once
+    let logup_hints = build_logup_hint_registry::<CircuitField<C>>();
+
+    // Use it for the frontend debug evaluation
+    debug_eval(&circuit, &assignment, logup_hints.clone());
+
+    // And reuse it for the witness solver
+    let witness = witness_solver
+        .solve_witness_with_hints(&assignment, &logup_hints)

Minor, but keeps intent and implementation aligned.

python/tests/circuit_parent_classes/test_ort_custom_layers.py (1)

92-93: Path consistency: use str(out_path) for the quantized model path.

After applying the fix above, update line 93 to use the same path variable:

     outputs_true = converter.run_model_onnx_runtime(path, inputs)
-    outputs_quant = converter.run_model_onnx_runtime("model.onnx", inputs)
+    outputs_quant = converter.run_model_onnx_runtime(str(out_path), inputs)

rust/jstprove_circuits/src/circuit_functions/layers/clip.rs (1)

106-109: Dead code: (None, None) match arm is unreachable.

The early return at lines 107-109 guarantees that at least one of min_input or max_input is Some. The (None, None) match arm at lines 191-198 can never execute.

Consider removing the dead arm or replacing with unreachable!() for clarity:

         match (min_bc_opt.as_ref(), max_bc_opt.as_ref()) {
-            (None, None) => {
-                // This case should have early-returned, but keep for completeness.
-                for &x_val in &x_bc {
-                    let clipped =
-                        constrained_clip(api, &range_ctx, &mut logup_ctx, x_val, None, None)?;
-                    out_storage.push(clipped);
-                }
-            }
             (Some(min_bc), None) => {

Also applies to: 190-198

rust/jstprove_circuits/src/circuit_functions/layers/maxpool.rs (1)

125-125: Use checked_sub for consistency with other layers.

Other layers (ReluLayer at line 56-59, ClipLayer at line 283-290) use checked_sub to derive shift_exponent. This direct subtraction could panic if n_bits is 0.

-            shift_exponent: layer_context.n_bits - 1,
+            shift_exponent: layer_context
+                .n_bits
+                .checked_sub(1)
+                .ok_or_else(|| LayerError::Other {
+                    layer: LayerKind::MaxPool,
+                    msg: "layer_context.n_bits must be at least 1".to_string(),
+                })?,

python/core/model_processing/onnx_quantizer/layers/base.py (2)

461-514: Verify attribute access on self relies on BaseOpQuantizer being present.

The quantize method calls self.add_scaled_initializer_inputs, self.get_scaling, and accesses self.new_initializers. These are provided by BaseOpQuantizer. The docstring warns about this, but there's no runtime check. Consider adding a defensive assertion or using typing.Protocol to make the contract explicit and catch errors earlier.

     def quantize(
         self,
         node: onnx.NodeProto,
         graph: onnx.GraphProto,
         scale_config: ScaleConfig,
         initializer_map: dict[str, onnx.TensorProto],
     ) -> list[onnx.NodeProto]:
         """Generic quantization template for most Int64 ops."""
+        # Defensive check: ensure BaseOpQuantizer is in the MRO
+        if not hasattr(self, "new_initializers"):
+            raise TypeError(
+                f"{self.__class__.__name__} must inherit from BaseOpQuantizer"
+            )
         _ = graph
         nodes = []

---

516-581: In-place graph modification in pre_analysis_transform may cause subtle issues.

The method modifies graph.initializer in place by deleting and appending items during iteration (lines 575-579). While iterating over scale_plan (not graph.initializer), the direct deletion pattern could be fragile if called multiple times or if the same tensor appears in multiple scale plan entries.

-            # Modify graph initializer in place
-            for j in range(len(graph.initializer)):
-                if graph.initializer[j].name == tensor.name:
-                    del graph.initializer[j]
-                    break
-            graph.initializer.append(new_tensor)
+            # Remove old initializer (if present) and append updated one
+            graph.initializer[:] = [
+                init for init in graph.initializer if init.name != tensor.name
+            ]
+            graph.initializer.append(new_tensor)

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (1)

99-109: Redundant shape validation after broadcast.

The broadcast_two_arrays helper (per the relevant snippet at tensor_ops.rs:222-237) already ensures both arrays have identical shapes. The length check on lines 99-109 is defensive but redundant. Consider removing it or converting to a debug assertion.

         let shape = a_bc.shape().to_vec();
-        if a_bc.len() != b_bc.len() {
-            return Err(LayerError::InvalidShape {
-                layer: LayerKind::Max,
-                msg: format!(
-                    "broadcast_two_arrays returned arrays of different sizes: {:?} vs {:?}",
-                    a_bc.shape(),
-                    b_bc.shape()
-                ),
-            }
-            .into());
-        }
+        debug_assert_eq!(
+            a_bc.len(),
+            b_bc.len(),
+            "broadcast_two_arrays should return same-sized arrays"
+        );

python/core/model_processing/onnx_quantizer/layers/clip.py (1)

80-92: check_supported returns None but base class signature returns str | None.

The method body does nothing and implicitly returns None, which matches the return type. However, adding an explicit return None would improve clarity and consistency with the base class pattern.

     def check_supported(
         self,
         node: onnx.NodeProto,
         initializer_map: dict[str, onnx.TensorProto] | None = None,
     ) -> None:
         """
         Minimal support check for Clip:

         - Clip is variadic elementwise with optional min/max as inputs or attrs.
         - We accept both forms; if attrs are present, ORT enforces semantics.
         - Broadcasting is ONNX-standard; we don't restrict further here.
         """
         _ = node, initializer_map
+        return None

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (2)

100-110: Same redundant shape validation as MaxLayer.

As with MaxLayer, the length check after broadcast_two_arrays is defensive but redundant since the helper guarantees identical shapes.

         let shape = a_bc.shape().to_vec();
-        if a_bc.len() != b_bc.len() {
-            return Err(LayerError::InvalidShape {
-                layer: LayerKind::Min,
-                msg: format!(
-                    "broadcast_two_arrays returned arrays of different sizes: {:?} vs {:?}",
-                    a_bc.shape(),
-                    b_bc.shape()
-                ),
-            }
-            .into());
-        }
+        debug_assert_eq!(
+            a_bc.len(),
+            b_bc.len(),
+            "broadcast_two_arrays should return same-sized arrays"
+        );

---

33-45: Consider extracting common layer struct pattern.

Both MaxLayer and MinLayer share identical struct fields. Consider extracting a common struct or macro to reduce duplication if more similar layers are planned.

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (1)

192-202: Minor inconsistency in error handling style.

new() panics on invalid chunk_bits while range_check() returns Result<(), CircuitError> for similar validation failures. Consider returning Result<Self, CircuitError> from new() for consistency, or document why panicking is preferred here (e.g., configuration errors are programmer mistakes vs. runtime input validation).

Current behavior is acceptable since it's documented, but unifying the error handling approach would improve API consistency.

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (1)

64-86: Hardcoded LayerKind::MaxPool in error messages for a generic context.

ShiftRangeContext is documented as shared across MaxPool, MaxLayer, MinLayer, and Clip (lines 29-32), but error messages on lines 71 and 76 hardcode LayerKind::MaxPool. This could be confusing when the error originates from a Min or Clip layer.

Consider parameterizing the layer kind or using a neutral identifier:

     pub fn new<C: Config, Builder: RootAPI<C>>(
         api: &mut Builder,
         shift_exponent: usize,
+        layer: LayerKind,
     ) -> Result<Self, LayerError> {
         let offset_: u32 = 1u32
             .checked_shl(
                 u32::try_from(shift_exponent).map_err(|_| LayerError::Other {
-                    layer: LayerKind::MaxPool,
+                    layer,
                     msg: format!("Shift exponent {shift_exponent} is too large for type: u32"),
                 })?,
             )
             .ok_or_else(|| LayerError::InvalidParameterValue {
-                layer: LayerKind::MaxPool,
+                layer,
                 layer_name: "ShiftRangeContext".to_string(),

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 14d7c15 and c8ae9cf.

⛔ Files ignored due to path filters (1)

• poetry.lock is excluded by !**/*.lock

📒 Files selected for processing (45)

• .github/PULL_REQUEST_TEMPLATE.md (2 hunks)
• python/core/circuits/base.py (5 hunks)
• python/core/circuits/errors.py (1 hunks)
• python/core/model_processing/converters/base.py (2 hunks)
• python/core/model_processing/onnx_custom_ops/__init__.py (1 hunks)
• python/core/model_processing/onnx_quantizer/exceptions.py (2 hunks)
• python/core/model_processing/onnx_quantizer/layers/base.py (1 hunks)
• python/core/model_processing/onnx_quantizer/layers/clip.py (1 hunks)
• python/core/model_processing/onnx_quantizer/layers/max.py (1 hunks)
• python/core/model_processing/onnx_quantizer/layers/min.py (1 hunks)
• python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py (3 hunks)
• python/core/model_templates/circuit_template.py (1 hunks)
• python/core/utils/errors.py (1 hunks)
• python/core/utils/scratch_tests.py (1 hunks)
• python/models/inputs/lenet_input.json (1 hunks)
• python/tests/circuit_e2e_tests/circuit_model_developer_test.py (13 hunks)
• python/tests/circuit_e2e_tests/helper_fns_for_tests.py (7 hunks)
• python/tests/circuit_parent_classes/test_ort_custom_layers.py (1 hunks)
• python/tests/onnx_quantizer_tests/layers/base.py (1 hunks)
• python/tests/onnx_quantizer_tests/layers/clip_config.py (1 hunks)
• python/tests/onnx_quantizer_tests/layers/max_config.py (1 hunks)
• python/tests/onnx_quantizer_tests/layers/min_config.py (1 hunks)
• python/tests/onnx_quantizer_tests/layers_tests/test_integration.py (1 hunks)
• python/tests/onnx_quantizer_tests/layers_tests/test_quantize.py (2 hunks)
• python/tests/onnx_quantizer_tests/test_registered_quantizers.py (4 hunks)
• rust-toolchain.toml (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/gadgets/mod.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/hints/bits.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/hints/max_min_clip.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/hints/mod.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/clip.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/layer_kinds.rs (2 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/max.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/maxpool.rs (4 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/min.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/mod.rs (2 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/relu.rs (6 hunks)
• rust/jstprove_circuits/src/circuit_functions/mod.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/utils/core_math.rs (0 hunks)
• rust/jstprove_circuits/src/circuit_functions/utils/mod.rs (0 hunks)
• rust/jstprove_circuits/src/circuit_functions/utils/quantization.rs (5 hunks)
• rust/jstprove_circuits/src/lib.rs (1 hunks)
• rust/jstprove_circuits/src/runner/main_runner.rs (4 hunks)

💤 Files with no reviewable changes (2)

• rust/jstprove_circuits/src/circuit_functions/utils/mod.rs
• rust/jstprove_circuits/src/circuit_functions/utils/core_math.rs

🧰 Additional context used

🧬 Code graph analysis (20)

rust/jstprove_circuits/src/circuit_functions/hints/max_min_clip.rs (1)

python/core/circuits/errors.py (1)

• CircuitError (6-32)

rust/jstprove_circuits/src/circuit_functions/gadgets/mod.rs (4)

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (4)

• range_check (233-256)
• constrained_reconstruct_from_bits (56-83)
• logup_range_check_pow2_unsigned (281-293)
• range_check_pow2_unsigned (122-137)

rust/jstprove_circuits/src/circuit_functions/layers/relu.rs (2)

• ShiftRangeContext (125-125)
• constrained_max (137-137)

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (3)

• constrained_clip (398-424)
• constrained_max (165-222)
• constrained_min (285-344)

rust/jstprove_circuits/src/circuit_functions/utils/quantization.rs (1)

• constrained_max (273-273)

rust/jstprove_circuits/src/circuit_functions/hints/mod.rs (3)

rust/jstprove_circuits/src/circuit_functions/hints/bits.rs (1)

• unconstrained_to_bits (74-105)

rust/jstprove_circuits/src/circuit_functions/hints/max_min_clip.rs (3)

• unconstrained_clip (241-261)
• unconstrained_max (72-99)
• unconstrained_min (155-182)

rust/jstprove_circuits/src/runner/main_runner.rs (3)

• build_logup_hint_registry (180-180)
• build_logup_hint_registry (292-292)
• build_logup_hint_registry (298-298)

python/tests/circuit_e2e_tests/circuit_model_developer_test.py (1)

python/tests/circuit_e2e_tests/helper_fns_for_tests.py (1)

• model_fixture (26-65)

python/tests/onnx_quantizer_tests/layers_tests/test_quantize.py (1)

python/tests/onnx_quantizer_tests/layers/base.py (2)

• LayerTestConfig (58-166)
• LayerTestSpec (27-55)

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (4)

rust/jstprove_circuits/src/circuit_functions/utils/onnx_model.rs (3)

• extract_params_and_expected_shape (168-200)
• get_input_name (247-259)
• get_optional_w_or_b (123-132)

rust/jstprove_circuits/src/circuit_functions/utils/tensor_ops.rs (2)

• broadcast_two_arrays (223-238)
• load_array_constants_or_get_inputs (62-81)

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (2)

• constrained_min (285-344)
• new (64-86)

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (1)

• new_default (208-210)

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (3)

rust/jstprove_circuits/src/circuit_functions/utils/onnx_model.rs (3)

• extract_params_and_expected_shape (168-200)
• get_input_name (247-259)
• get_optional_w_or_b (123-132)

rust/jstprove_circuits/src/circuit_functions/utils/tensor_ops.rs (2)

• broadcast_two_arrays (223-238)
• load_array_constants_or_get_inputs (62-81)

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (3)

• constrained_max (165-222)
• logup_ctx (208-209)
• logup_ctx (330-331)

python/core/model_processing/onnx_quantizer/layers/max.py (3)

python/core/model_processing/onnx_quantizer/layers/base.py (3)

• BaseOpQuantizer (29-417)
• QuantizerBase (420-581)
• ScaleConfig (23-26)

python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py (1)

• quantize (120-159)

python/core/model_processing/onnx_quantizer/layers/min.py (2)

• quantize (36-44)
• check_supported (46-53)

rust/jstprove_circuits/src/circuit_functions/layers/maxpool.rs (5)

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (4)

• constrained_max (165-222)
• new (64-86)
• logup_ctx (208-209)
• logup_ctx (330-331)

rust/jstprove_circuits/src/circuit_functions/layers/clip.rs (2)

• logup_ctx (185-185)
• logup_ctx (242-242)

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (2)

• logup_ctx (88-88)
• logup_ctx (120-120)

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (2)

• logup_ctx (88-88)
• logup_ctx (121-121)

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (1)

• new_default (208-210)

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (2)

rust/jstprove_circuits/src/circuit_functions/hints/max_min_clip.rs (2)

• unconstrained_max (72-99)
• unconstrained_min (155-182)

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (1)

• new (192-202)

rust/jstprove_circuits/src/circuit_functions/layers/relu.rs (4)

rust/jstprove_circuits/src/circuit_functions/utils/quantization.rs (6)

• constrained_max (273-273)
• new (116-161)
• logup_ctx (235-236)
• logup_ctx (252-253)
• logup_ctx (335-335)
• logup_ctx (345-345)

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (2)

• new (192-202)
• new_default (208-210)

rust/jstprove_circuits/src/circuit_functions/layers/clip.rs (2)

• logup_ctx (185-185)
• logup_ctx (242-242)

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (2)

• logup_ctx (88-88)
• logup_ctx (120-120)

python/tests/onnx_quantizer_tests/layers/min_config.py (1)

python/tests/onnx_quantizer_tests/layers/base.py (10)

• e2e_test (247-248)
• edge_case_test (243-244)
• valid_test (235-236)
• BaseLayerConfigProvider (251-277)
• LayerTestConfig (58-166)
• description (175-177)
• override_input_shapes (195-197)
• tags (215-217)
• override_initializer (187-189)
• skip (219-221)

python/tests/circuit_parent_classes/test_ort_custom_layers.py (1)

python/core/model_processing/converters/onnx_converter.py (1)

• run_model_onnx_runtime (301-346)

rust/jstprove_circuits/src/runner/main_runner.rs (1)

rust/jstprove_circuits/src/circuit_functions/hints/mod.rs (1)

• build_logup_hint_registry (16-24)

rust/jstprove_circuits/src/circuit_functions/utils/quantization.rs (5)

rust/jstprove_circuits/src/circuit_functions/layers/relu.rs (4)

• ShiftRangeContext (125-125)
• constrained_max (137-137)
• logup_ctx (129-129)
• logup_ctx (142-142)

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (4)

• constrained_max (165-222)
• logup_ctx (208-209)
• logup_ctx (330-331)
• new (64-86)

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (2)

• logup_ctx (88-88)
• logup_ctx (120-120)

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (2)

• logup_ctx (88-88)
• logup_ctx (121-121)

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (2)

• new (192-202)
• new_default (208-210)

python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py (3)

python/core/model_processing/onnx_quantizer/layers/clip.py (1)

• ClipQuantizer (49-92)

python/core/model_processing/onnx_quantizer/layers/max.py (1)

• MaxQuantizer (24-47)

python/core/model_processing/onnx_quantizer/layers/min.py (1)

• MinQuantizer (24-53)

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (1)

rust/jstprove_circuits/src/circuit_functions/hints/bits.rs (1)

• unconstrained_to_bits (74-105)

python/core/circuits/base.py (2)

python/core/circuits/errors.py (3)

• CircuitConfigurationError (35-57)
• CircuitInputError (60-90)
• CircuitProcessingError (144-167)

python/core/utils/errors.py (1)

• ShapeMismatchError (104-117)

python/tests/onnx_quantizer_tests/layers/max_config.py (1)

python/tests/onnx_quantizer_tests/layers/base.py (11)

• e2e_test (247-248)
• edge_case_test (243-244)
• valid_test (235-236)
• BaseLayerConfigProvider (251-277)
• LayerTestConfig (58-166)
• description (175-177)
• override_input_shapes (195-197)
• tags (215-217)
• override_initializer (187-189)
• override_output_shapes (199-201)
• skip (219-221)

rust/jstprove_circuits/src/circuit_functions/layers/layer_kinds.rs (5)

rust/jstprove_circuits/src/circuit_functions/layers/clip.rs (1)

• build (254-305)

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (1)

• build (132-173)

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (1)

• build (133-174)

rust/jstprove_circuits/src/circuit_functions/layers/add.rs (1)

• build (71-99)

rust/jstprove_circuits/src/circuit_functions/layers/layer_ops.rs (1)

• build (51-60)

🪛 Ruff (0.14.7)

python/tests/circuit_e2e_tests/circuit_model_developer_test.py

43-43: Redefinition of unused model_fixture from line 30

(F811)

python/core/model_processing/onnx_quantizer/exceptions.py

154-154: Unused noqa directive (non-enabled: FBT001)

Remove unused noqa directive

(RUF100)

python/core/model_processing/onnx_custom_ops/__init__.py

16-16: Unused noqa directive (non-enabled: PYI056)

Remove unused noqa directive

(RUF100)

python/core/utils/scratch_tests.py

26-26: Unused noqa directive (non-enabled: T201)

Remove unused noqa directive

(RUF100)

---

59-59: Unused noqa directive (non-enabled: T201)

Remove unused noqa directive

(RUF100)

---

64-64: Unused noqa directive (non-enabled: T201)

Remove unused noqa directive

(RUF100)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: End-to-End Testing (3.10)
• GitHub Check: End-to-End Testing (3.11)
• GitHub Check: End-to-End Testing (3.12)
• GitHub Check: Check Formatting and Linting

🔇 Additional comments (55)

.github/PULL_REQUEST_TEMPLATE.md (1)

6-6: LGTM! Consistent placeholder formatting cleanup.

The formatting changes normalize the placeholder lines across all three sections ("Related Issue," "Deployment Notes," "Additional Comments") from -- to -, which is a standard lint cleanup. No structural or functional impact to the template.

Also applies to: 27-27, 31-31

rust-toolchain.toml (1)

2-2: Valid TOML syntax restored.

The leading dash before channel was invalid TOML syntax and would cause parse errors during build/CI setup. The fix correctly restores the standard key-value pair format.

python/models/inputs/lenet_input.json (1)

3078-3078: Formatting change approved.

The trailing newline addition aligns with the pre-commit formatting pass mentioned in the PR objectives. JSON structure and data remain unchanged.

python/core/utils/errors.py (1)

33-40: Formatting cleanup looks good.

The ProofBackendError.__init__ method signature has been cleanly reformatted with proper indentation and type hints. All parameters and logic remain intact.

python/core/model_processing/onnx_custom_ops/__init__.py (1)

3-11: LGTM! Good modernization with pathlib and type hints.

The migration from os.path to pathlib.Path and the addition of type annotations follow modern Python best practices.

python/core/utils/scratch_tests.py (3)

1-5: Good modernization of imports and type hints setup.

The addition of from __future__ import annotations enables cleaner type hint syntax throughout the file, and the consolidated ONNX imports are appropriate.

---

8-13: Excellent addition of type hints and docstring.

The explicit type annotations and concise docstring improve code clarity and maintainability.

---

29-34: Well-structured improvements to function signature and error handling.

The type hints, docstring, and error message pattern (storing in a variable before raising) follow Python best practices.

Also applies to: 50-51

python/core/model_templates/circuit_template.py (3)

1-4: Imports and randomness choice look good

Using from __future__ import annotations aligns with the modern type hints (dict[str, int], str | None), and switching nonce generation to secrets.randbelow is safe and appropriate for a nonce-style value.

---

26-31: Constructor/file_name wiring and nonce default

Storing file_name on self and keeping the __init__ fully backward compatible (only adding an optional arg) looks good. One minor behavioral change is the nonce range: randbelow(10_000) yields [0, 9999]. If any downstream logic or tests assume a 5‑digit upper bound including 10000 or a non‑zero nonce, it’s worth double‑checking those expectations.

Also applies to: 37-40

---

42-51: get_inputs implementation is clear and typed

The explicit dict[str, int] return with named keys ("input_a", "input_b", "nonce") is straightforward and matches required_keys. No functional issues here.

python/tests/circuit_e2e_tests/circuit_model_developer_test.py (1)

5-9: LGTM: TYPE_CHECKING guard avoids runtime import overhead.

Moving Generator inside the TYPE_CHECKING block is the correct pattern to avoid importing it at runtime when it's only needed for type annotations.

python/core/model_processing/onnx_quantizer/exceptions.py (1)

34-41: LGTM: Signature formatting improves readability.

The multiline signature format makes the parameter list easier to scan and maintains consistency with other exception classes in the file.

python/core/circuits/errors.py (1)

70-77: LGTM: Signature formatting improves consistency.

The explicit multiline signature with -> None return type annotation improves consistency with other exception classes in the file.

python/tests/circuit_e2e_tests/helper_fns_for_tests.py (2)

3-5: LGTM: Type annotation modernization to Python 3.10+ syntax.

The changes modernize type annotations using native union syntax (|) and proper TYPE_CHECKING guards, improving readability and following current Python best practices.

Also applies to: 112-115, 138-138

---

68-68: LGTM: Fixture decorator cleanup.

Removing parentheses from @pytest.fixture() to @pytest.fixture is consistent with the cleanup in circuit_model_developer_test.py.

Also applies to: 79-79, 90-90, 101-101

python/core/circuits/base.py (3)

7-7: LGTM: NumPy import modernization.

Changing from from numpy import asarray, ndarray to import numpy as np is a standard best practice that improves readability and namespace clarity.

---

803-815: LGTM: Enhanced shape handling for dict inputs.

The new logic properly extracts per-key shapes from input_shape when it's a dict, with appropriate error handling for missing keys. The explicit shape type validation improves error messages.

Also applies to: 820-824

---

827-829: LGTM: Consistent use of np.asarray throughout.

All conversions now consistently use np.asarray() with the np. prefix, maintaining consistency with the updated import style.

Also applies to: 844-846, 887-897

python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py (1)

23-23: LGTM: New quantizers integrated correctly.

The Clip, Max, and Min quantizers are imported and registered following the same pattern as existing quantizers (Add, Sub, Mul, etc.), with proper new_initializers sharing for tracking newly created constants.

Also applies to: 29-29, 31-31, 80-80, 90-91

python/tests/onnx_quantizer_tests/layers/base.py (1)

116-116: LGTM: Documentation cleanup.

Merging the two comments improves readability without changing functionality.

python/core/model_processing/converters/base.py (1)

5-5: LGTM: Type alias modernization to Python 3.10+ syntax.

Removing Optional and Union in favor of native | syntax and None union improves readability and follows current Python typing conventions.

Also applies to: 17-20, 22-22

python/tests/onnx_quantizer_tests/layers_tests/test_integration.py (1)

162-177: LGTM! Improved naming clarity.

The renaming from scaled_inputs to quantized_inputs better reflects that these are inputs for the quantized model, and the updated comment accurately describes the float64 casting for ORT compatibility.

rust/jstprove_circuits/src/circuit_functions/layers/mod.rs (1)

3-3: LGTM! New layer modules properly exposed.

The three new modules (clip, max, min) are correctly declared as public, expanding the API surface to support the new Clip/Max/Min layer implementations.

Also applies to: 12-12, 14-14

rust/jstprove_circuits/src/circuit_functions/mod.rs (1)

2-3: LGTM! Gadgets and hints modules properly exposed.

The new gadgets and hints modules expand the public API to expose the underlying constraint gadgets and witness helpers that support the new layer implementations.

python/tests/onnx_quantizer_tests/test_registered_quantizers.py (2)

124-129: LGTM! More flexible input validation.

The conditional assertion correctly handles operators that may not require inputs (like Constant), checking for inputs only when they're actually required based on inputs list.

---

78-85: Opset versions are appropriate for all tested operators.

Min, Max, and Clip are standard ONNX operators introduced in opset 1 with current specification at opset 13. Opset 22 for the default domain fully supports all three operators. The ai.onnx.contrib opset 1 is correctly specified for any custom operators that may use that domain.

rust/jstprove_circuits/src/lib.rs (2)

1-27: LGTM! Improved documentation and appropriate lint suppressions.

The documentation updates better describe the crate structure and module purposes. The clippy allow attributes for doc-related lints are appropriate for technical documentation with specific formatting requirements.

---

28-28: Use min_specialization without concerns.

The change from specialization to min_specialization is safe and correct. Verification confirms the crate contains no specialization usage patterns (no default impl, no overlapping trait implementations, no specialized generic bounds). Since min_specialization is a strict subset designed for exactly this use case, the crate will compile and function correctly with this feature flag. This change improves forward compatibility as min_specialization is closer to stabilization than full specialization.

python/tests/onnx_quantizer_tests/layers/min_config.py (1)

18-38: LGTM! Well-structured Min layer configuration.

The MinConfigProvider follows the established pattern for layer test configurations. The configuration appropriately defines the Min operation with two inputs (A, B) and matching shapes.

python/tests/onnx_quantizer_tests/layers_tests/test_quantize.py (1)

40-43: LGTM! Clean normalization of return type.

The change to consistently return list[onnx.NodeProto] improves the interface clarity. The normalization logic at lines 71-72 maintains backward compatibility by wrapping single-node results in a list.

Also applies to: 71-72

rust/jstprove_circuits/src/circuit_functions/layers/layer_kinds.rs (2)

12-12: LGTM! New layer imports properly added.

The imports for ClipLayer, MaxLayer, and MinLayer are correctly added to support the new layer registrations.

Also applies to: 17-17, 19-19

---

133-146: LGTM! Layer registry properly updated with new layers.

The three new layer types (Clip, Max, Min) are correctly registered with their respective builder functions. The reformatting improves readability and maintains consistency across the registry entries.

python/tests/onnx_quantizer_tests/layers/clip_config.py (1)

17-107: ClipConfigProvider looks consistent with the shared test harness

Layer config and test specs (scalar, broadcast, initializer, e2e, and empty-tensor cases) are coherent with LayerTestConfig’s input/initializer semantics and ONNX Clip behavior. Nothing blocking here.

rust/jstprove_circuits/src/circuit_functions/gadgets/mod.rs (1)

5-12: Gadget re-exports are clean and coherent

The module layout and re-exports (max/min/clip + LogUp range-check helpers) provide a clear, centralized gadget API for layers without leaking internal structure. Looks good.

rust/jstprove_circuits/src/circuit_functions/layers/clip.rs (2)

120-142: Broadcasting logic is correct.

The pairwise broadcasting approach properly handles the case where min is broadcast first, then X expands due to max, requiring min to be re-broadcast. The shape validation at lines 146-171 provides good defense-in-depth.

---

269-280: LGTM!

The optional initializer handling for X, min, and max correctly follows the pattern used by other layers. This supports both runtime inputs and constant-folded initializers.

rust/jstprove_circuits/src/circuit_functions/layers/relu.rs (2)

119-146: LGTM!

The relu_array function correctly implements ReLU as max(x, 0) using the shared gadget infrastructure. Key observations:

• Zero constant is created once outside the loop (line 131)
• LogUp context is properly scoped: initialized before processing, finalized after
• Error handling correctly propagates through the collect() pattern

---

46-63: LGTM!

The apply method cleanly derives shift_exponent from n_bits with proper bounds checking, then delegates to relu_array. The switch to .as_str() for HashMap lookup is a minor efficiency improvement.

rust/jstprove_circuits/src/circuit_functions/layers/maxpool.rs (2)

431-434: LGTM!

The LogUp context lifecycle is correctly implemented:

• ShiftRangeContext and LogupRangeCheckContext are created before the pooling loops
• All constrained_max calls share the same context
• finalize is called once after all operations complete

This matches the pattern used in ClipLayer, MaxLayer, MinLayer, and ReluLayer.

Also applies to: 474-475

---

156-230: LGTM!

The setup_maxpooling_2d_params function provides comprehensive normalization for pooling parameters with clear error messages for invalid input lengths.

rust/jstprove_circuits/src/circuit_functions/hints/mod.rs (1)

14-24: LGTM!

The hint registry builder is well-structured and properly used in main_runner.rs. However, note that the comment referencing new_hint(...) calls cannot be verified as no such calls exist in the codebase—this may be forward-looking documentation or refer to hint identifiers used in a different context (e.g., circuit definitions outside Rust code).

python/core/model_processing/onnx_quantizer/layers/base.py (1)

420-459: Well-documented mixin with clear integration requirements.

The docstring clearly explains that QuantizerBase must be combined with BaseOpQuantizer via multiple inheritance and lists the required methods/attributes. This helps prevent misuse at runtime.

rust/jstprove_circuits/src/circuit_functions/utils/quantization.rs (2)

214-283: LogUp-based rescaling implementation looks correct.

The function properly:

1. Computes the shifted dividend and performs unconstrained division
2. Enforces the equality constraint
3. Uses LogUp for range checks on remainder and quotient
4. Correctly constructs ShiftRangeContext for the ReLU branch
5. Applies constrained_max for ReLU

The error wrapping uses descriptive messages that include the underlying error context.

---

324-349: Proper lifecycle management of LogupRangeCheckContext.

The rescale_array function correctly:

1. Creates a shared LogupRangeCheckContext with new_default()
2. Initializes it before use with init()
3. Applies rescale to each element using the shared context
4. Finalizes with finalize() after all elements are processed

This ensures the LogUp consistency check covers all queries in the batch.

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (1)

49-130: MaxLayer apply implementation follows established patterns correctly.

The implementation mirrors AddLayer's structure for input resolution, broadcasting, and output construction. The LogUp context is properly initialized before use and finalized after all elements are processed.

python/core/model_processing/onnx_quantizer/layers/clip.py (1)

15-46: Well-documented QuantizeClip mixin with clear semantics.

The docstring explains why USE_WB = True is appropriate (min/max need the same fixed-point scale as X) and why USE_SCALING = False (Clip doesn't introduce its own scale). The SCALE_PLAN correctly maps inputs 1 and 2 to 1x scaling.

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (1)

49-131: MinLayer apply implementation is consistent with MaxLayer.

The implementation correctly mirrors MaxLayer's structure, using constrained_min instead of constrained_max. LogUp context lifecycle (init/finalize) is properly managed.

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (4)

1-17: LGTM!

Module documentation is clear and imports are well-organized with appropriate section comments.

---

56-83: LGTM - logic is sound.

The bit reconstruction and boolean validation are correctly implemented with proper overflow checks via checked_shl.

Minor observation: the error on line 73-76 for checked_shl failure reports max: u128::from(u32::MAX), but the actual constraint is i <= 31 (max shift for u32). This doesn't affect functionality but could be slightly clearer for debugging.

---

122-137: LGTM!

Clean composition of bit decomposition, reconstruction, and equality assertion. Error propagation is handled correctly via the ? operator.

---

281-293: LGTM!

Clean one-shot helper that properly manages the LogUp context lifecycle. Error propagation is correct.

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (3)

165-222: LGTM - well-designed constrained max gadget.

The algorithm is sound:

1. Shift to nonnegative space to avoid signed arithmetic issues
2. Compute unconstrained max as a witness hint
3. Range-check all deltas M - x_i ≥ 0 via LogUp
4. Assert ∏ (M - x_i) = 0 to enforce that M equals at least one input

This correctly enforces the max constraint under the stated assumptions. Documentation is thorough and helpful.

---

285-344: LGTM!

Correctly mirrors constrained_max with the appropriate delta reversal (x_i - M instead of M - x_i). The symmetry is well-maintained.

---

398-424: LGTM!

Clean composition of max/min gadgets matching ONNX Clip semantics. All four bound combinations are handled correctly:

• Both bounds: min(max(x, lower), upper)
• Lower only: max(x, lower)
• Upper only: min(x, upper)
• Neither: identity

[coderabbitai]

Actionable comments posted: 1

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a1a216d and 148872f.

📒 Files selected for processing (1)

• python/tests/onnx_quantizer_tests/layers/max_config.py (1 hunks)

🧰 Additional context used

🧠 Learnings (1)

📓 Common learnings

Learnt from: tmfreiberg
Repo: inference-labs-inc/JSTprove PR: 85
File: python/core/model_processing/onnx_quantizer/layers/clip.py:59-65
Timestamp: 2025-12-05T18:46:45.499Z
Learning: In python/core/model_processing/onnx_quantizer/layers/, newer quantizers (ClipQuantizer, MaxQuantizer, MinQuantizer) use a shared dict pattern for new_initializers (dict[str, onnx.TensorProto] | None) and directly assign self.new_initializers without calling super().__init__(), intentionally bypassing BaseOpQuantizer.__init__ which would initialize it as a list.

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)

• GitHub Check: Check Formatting and Linting
• GitHub Check: End-to-End Testing (3.12)
• GitHub Check: End-to-End Testing (3.11)

🔇 Additional comments (3)

python/tests/onnx_quantizer_tests/layers/max_config.py (3)

18-38: LGTM!

The MaxConfigProvider class and get_config method are well-structured. The configuration correctly defines the Max operator with two dynamic inputs, no attributes (appropriate for Max), and consistent input/output shapes.

---

74-79: Previous review feedback addressed.

The empty_tensor test now correctly includes .override_output_shapes(max_output=[0]) to align with the empty input shapes [0].

---

40-86: Well-structured test specifications.

The test suite provides good coverage across:

• Basic elementwise operations
• Broadcasting scenarios
• Initializer handling
• End-to-end flows
• Edge cases (empty tensors)
• Performance/stress tests (appropriately skipped by default)

The seeded RNG ensures reproducibility.

[coderabbitai]

Actionable comments posted: 1

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 148872f and f223b64.

📒 Files selected for processing (1)

• rust/jstprove_circuits/src/circuit_functions/hints/bits.rs (1 hunks)

🧰 Additional context used

🧠 Learnings (1)

📓 Common learnings

Learnt from: tmfreiberg
Repo: inference-labs-inc/JSTprove PR: 85
File: python/core/model_processing/onnx_quantizer/layers/clip.py:59-65
Timestamp: 2025-12-05T18:46:45.499Z
Learning: In python/core/model_processing/onnx_quantizer/layers/, newer quantizers (ClipQuantizer, MaxQuantizer, MinQuantizer) use a shared dict pattern for new_initializers (dict[str, onnx.TensorProto] | None) and directly assign self.new_initializers without calling super().__init__(), intentionally bypassing BaseOpQuantizer.__init__ which would initialize it as a list.

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: End-to-End Testing (3.10)
• GitHub Check: End-to-End Testing (3.11)
• GitHub Check: Check Formatting and Linting
• GitHub Check: End-to-End Testing (3.12)

🔇 Additional comments (4)

rust/jstprove_circuits/src/circuit_functions/hints/bits.rs (4)

1-16: LGTM: Clear documentation and appropriate imports.

The warning that these helpers are witness-only and not soundness-critical is important context. The imports are minimal and all utilized.

---

22-73: Excellent documentation.

The function documentation is thorough, clearly explains the unconstrained nature of the operation, and provides a correct worked example. The warnings about lack of soundness are appropriate.

---

90-91: Error message fixed from prior review.

The error message is now complete and descriptive, addressing the past review comment. Well done.

---

95-104: Correct bit-extraction loop.

The iterative approach correctly extracts the least significant bit using unconstrained_bit_and, appends it, and shifts right. The result is in little-endian order as documented.

[coderabbitai]

Actionable comments posted: 1

♻️ Duplicate comments (1)

python/tests/circuit_parent_classes/test_ort_custom_layers.py (1)

93-93: Hardcoded path causes test failure — use out_path instead.

Line 93 still references the hardcoded "model.onnx" path instead of using out_path where the quantized model was actually saved (line 82). This is the direct cause of the pipeline failure: ONNX runtime could not load the model 'model.onnx' (NO_SUCHFILE).

Apply this diff to fix the issue:

-    outputs_quant = converter.run_model_onnx_runtime("model.onnx", inputs)
+    outputs_quant = converter.run_model_onnx_runtime(out_path, inputs)

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between f223b64 and 04fc017.

📒 Files selected for processing (6)

• python/core/model_processing/onnx_quantizer/layers/max.py (1 hunks)
• python/core/model_processing/onnx_quantizer/layers/min.py (1 hunks)
• python/tests/circuit_parent_classes/test_ort_custom_layers.py (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/hints/max_min_clip.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/max.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/min.rs (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• python/core/model_processing/onnx_quantizer/layers/max.py

🧰 Additional context used

🧬 Code graph analysis (1)

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (5)

rust/jstprove_circuits/src/circuit_functions/utils/onnx_model.rs (3)

• extract_params_and_expected_shape (168-200)
• get_input_name (247-259)
• get_optional_w_or_b (123-132)

rust/jstprove_circuits/src/circuit_functions/utils/tensor_ops.rs (2)

• broadcast_two_arrays (223-238)
• load_array_constants_or_get_inputs (62-81)

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (4)

• constrained_min (285-344)
• new (64-86)
• logup_ctx (208-209)
• logup_ctx (330-331)

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (4)

• apply (50-130)
• logup_ctx (88-88)
• logup_ctx (120-120)
• build (132-188)

rust/jstprove_circuits/src/circuit_functions/layers/clip.rs (4)

• apply (63-252)
• logup_ctx (185-185)
• logup_ctx (242-242)
• build (254-305)

🪛 GitHub Actions: Unit, Integration Python Tests

python/tests/circuit_parent_classes/test_ort_custom_layers.py

[error] 1-1: Integration test failure: ONNX runtime could not load the model 'model.onnx' (NO_SUCHFILE). The test 'test_tiny_conv' expects an existing ONNX model at this path; ensure the model is generated/saved before running this test. Command: pytest --integration --html=report/integration_tests.html

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

• GitHub Check: End-to-End Testing (3.11)
• GitHub Check: Check Formatting and Linting

🔇 Additional comments (7)

python/tests/circuit_parent_classes/test_ort_custom_layers.py (1)

12-58: LGTM!

The fixture correctly uses the tmp_path fixture to create an isolated temporary ONNX model. The snake_case variable naming (w_init, z_init) and proper Path return type are good improvements.

rust/jstprove_circuits/src/circuit_functions/hints/max_min_clip.rs (4)

1-10: LGTM! Clear module documentation and appropriate imports.

The module documentation clearly establishes that these helpers are witness-only and not sound on their own, with a helpful pointer to the constrained counterparts. The imports are minimal and well-organized.

---

72-99: LGTM! Well-designed witness computation with proper error handling.

The selector-based approach elegantly handles the max computation without branching, and the extensive documentation clearly explains the assumptions and limitations. The empty-slice error handling is correct and the error message was properly fixed in the previous commit.

---

155-182: LGTM! Correct mirror implementation of unconstrained_max.

The minimum computation correctly mirrors the max implementation with appropriate comparison operators. The error message properly references unconstrained_min (fixed in commit f4c5c31), and the documentation maintains the same high quality.

---

241-261: LGTM! Clean composition with proper optional bound handling.

The clip implementation elegantly composes unconstrained_max and unconstrained_min to handle optional bounds. The function never passes empty slices (always passes 2-element arrays), and the error propagation is correct. The documentation appropriately notes that constraint verification (e.g., lower <= upper) is deferred to the constrained layer.

python/core/model_processing/onnx_quantizer/layers/min.py (1)

15-54: MinQuantizer wiring and metadata look consistent with the shared QuantizerBase pattern

OP/DOMAIN metadata, passthrough scaling flags, initializer handling in __init__, and delegation in quantize all align with the Clip/Max quantizers and preserve the BaseOpQuantizer invariants on new_initializers. No changes requested here.

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (1)

31-189: MaxLayer implementation is robust and consistent with other elementwise layers

Input resolution, broadcasting, range-context setup, LogUp reuse, and error handling (including safe .get() on layer.inputs and checked shift_exponent) are all well-structured and mirror the Add/Clip patterns. No issues spotted.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (1)

133-189: Minor cleanups: unused input_shape, dead_code allowance, and comment wording

The build path and error mapping look correct and consistent with other layers, and the LayerKind::Min usage for both inputs is now coherent. A few optional nits you might consider:

• input_shape is populated from expected_shape but not used in apply. Either assert that the broadcast shape matches self.input_shape (for extra safety) or drop the field to avoid redundant state.
• With the type actively used via LayerOp, #[allow(dead_code)] on MinLayer may no longer be necessary.
• The comment “Match MaxPool’s choice: use n_bits - 1 as the shift exponent” might be misleading if the intended reference is actually the Max layer rather than a pooling layer; consider aligning the wording with the real precedent for clarity.

These are all non-blocking polish items; behaviorally the build logic looks fine.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 04fc017 and 8f0f7c4.

📒 Files selected for processing (3)

• python/tests/onnx_quantizer_tests/layers/max_config.py (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/hints/bits.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/min.rs (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• rust/jstprove_circuits/src/circuit_functions/hints/bits.rs

🧰 Additional context used

🧬 Code graph analysis (1)

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (3)

rust/jstprove_circuits/src/circuit_functions/utils/tensor_ops.rs (2)

• broadcast_two_arrays (223-238)
• load_array_constants_or_get_inputs (62-81)

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (4)

• constrained_min (285-344)
• new (64-86)
• logup_ctx (208-209)
• logup_ctx (330-331)

rust/jstprove_circuits/src/circuit_functions/gadgets/range_check.rs (1)

• new_default (208-210)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

• GitHub Check: build
• GitHub Check: Check Formatting and Linting

🔇 Additional comments (4)

python/tests/onnx_quantizer_tests/layers/max_config.py (3)

1-16: LGTM!

The imports are well-organized and appropriate for the test configuration provider. The use of from __future__ import annotations enables modern type hints.

---

18-23: LGTM!

The class definition and layer_name property follow the expected pattern for layer config providers in the test suite.

---

25-38: LGTM!

The configuration correctly describes the ONNX Max operator: no layer-specific attributes, two dynamic inputs by default, and aligned input/output shapes.

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (1)

49-131: Elementwise constrained_min wiring and LogUp context usage look solid

The apply implementation cleanly mirrors the Add/Max patterns: resolving input names, loading constants or runtime inputs, broadcasting with broadcast_two_arrays, building a per-layer ShiftRangeContext, and sharing a LogupRangeCheckContext across all pairwise constrained_min calls, with finalize at the end. The defensive len equality check after broadcasting and the reshape error mapping both look reasonable. I don’t see functional or safety issues here; this is a clear, idiomatic integration of the gadget into a layer op.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

python/tests/onnx_quantizer_tests/layers/max_config.py (1)

40-90: Good coverage of Max scenarios; consider tightening return type

The test spec set covers basic, broadcasting, initializer-based, e2e, empty-tensor, and large-tensor stress cases, with dtype and empty-output-shape details handled correctly. As a small follow-up, you could refine the signature of get_test_specs to return a more specific type (e.g., list[LayerTestSpec]) if that alias is available from the factory module, to help static type checkers.

python/tests/circuit_parent_classes/test_ort_custom_layers.py (1)

68-93: Inconsistent Path-to-string conversions for file operations.

The code mixes Path objects and strings when passing paths to ONNX/runtime functions:

• Line 68: onnx.load(path) passes a Path directly
• Line 87: onnx.load(str(out_path)) explicitly converts to string
• Line 92: run_model_onnx_runtime(path, ...) passes a Path directly
• Line 93: run_model_onnx_runtime(out_path, ...) passes a Path directly

The run_model_onnx_runtime signature expects path: str, not Path. While many libraries accept PathLike objects at runtime, explicit conversion maintains type correctness and avoids potential issues with type checkers or specific library versions.

Apply this diff for consistency:

     # Load and validate original model
-    model = onnx.load(path)
+    model = onnx.load(str(path))
     onnx.checker.check_model(model)

     # Prepare inputs and compare outputs
     inputs = np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4)
-    outputs_true = converter.run_model_onnx_runtime(path, inputs)
-    outputs_quant = converter.run_model_onnx_runtime(out_path, inputs)
+    outputs_true = converter.run_model_onnx_runtime(str(path), inputs)
+    outputs_quant = converter.run_model_onnx_runtime(str(out_path), inputs)

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 8f0f7c4 and 3028d4b.

📒 Files selected for processing (3)

• python/core/model_processing/onnx_custom_ops/__init__.py (1 hunks)
• python/tests/circuit_parent_classes/test_ort_custom_layers.py (1 hunks)
• python/tests/onnx_quantizer_tests/layers/max_config.py (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (2)

python/tests/onnx_quantizer_tests/layers/max_config.py (1)

python/tests/onnx_quantizer_tests/layers/base.py (6)

• e2e_test (247-248)
• edge_case_test (243-244)
• valid_test (235-236)
• BaseLayerConfigProvider (251-277)
• LayerTestConfig (58-166)
• skip (219-221)

python/tests/circuit_parent_classes/test_ort_custom_layers.py (4)

python/core/model_processing/converters/onnx_converter.py (2)

• ONNXConverter (121-1220)
• run_model_onnx_runtime (301-346)

python/tests/circuit_parent_classes/test_onnx_converter.py (1)

• converter (45-49)

python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py (1)

• check_model (161-180)

python/core/model_processing/converters/base.py (1)

• quantize_model (84-103)

🪛 Ruff (0.14.7)

python/core/model_processing/onnx_custom_ops/__init__.py

17-17: Unused noqa directive (non-enabled: PYI056)

Remove unused noqa directive

(RUF100)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: End-to-End Testing (3.10)
• GitHub Check: End-to-End Testing (3.11)
• GitHub Check: End-to-End Testing (3.12)
• GitHub Check: Check Formatting and Linting

🔇 Additional comments (6)

python/tests/onnx_quantizer_tests/layers/max_config.py (2)

4-15: Imports and RNG usage look consistent and deterministic

Using np.random.default_rng(TEST_RNG_SEED) plus .astype(np.float32) for initializers is aligned with ONNX float32 expectations and mirrors patterns in other layer configs. No issues here.

---

18-38: Config wiring for Max layer is clear and aligned with BaseLayerConfigProvider

layer_name="Max", op_type="Max", inputs ["A", "B"], and the symmetric [1, 3, 4, 4] shapes for both inputs and max_output match the intended elementwise Max semantics and the test harness expectations. This looks good.

python/core/model_processing/onnx_custom_ops/__init__.py (2)

3-3: LGTM: Modern pathlib usage.

The migration from os.path to pathlib.Path aligns with modern Python best practices.

Also applies to: 9-9

---

12-12: LGTM: Type annotation improves code quality.

Adding the list[str] type hint to __all__ enhances static type checking.

python/tests/circuit_parent_classes/test_ort_custom_layers.py (2)

12-58: LGTM! Fixture properly uses tmp_path for test isolation.

The fixture correctly accepts tmp_path, saves the model to an isolated location, and returns a Path object. The ONNX model construction is clean with consistent string literals.

---

90-98: LGTM! Test comparison logic is sound.

The test correctly compares original and quantized model outputs using the appropriate scaling factor (2^21 matching the quantization parameters) with reasonable tolerances for floating-point comparison.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

rust/jstprove_circuits/src/circuit_functions/layers/max.rs (1)

49-188: MaxLayer implementation correctly wires constrained_max and shared LogUp context

The MaxLayer apply/build flow looks solid: inputs are resolved via get_input_name and load_array_constants_or_get_inputs, broadcasted with broadcast_two_arrays, and then passed elementwise through constrained_max using a single ShiftRangeContext and shared LogupRangeCheckContext, with clear error paths (LayerError::Other/InvalidShape, MissingInput, and n_bits.checked_sub(1)). This matches the intended gadget‑based Max behavior and integrates cleanly with the existing layer utilities.

If you find Max/Min evolving further, consider factoring the common “broadcast + per‑element constrained_{max,min} + LogUp finalize” pattern into a small internal helper to keep the two layers in lockstep and avoid future divergence, but it’s not strictly necessary right now.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 3028d4b and 93a8246.

📒 Files selected for processing (4)

• python/tests/onnx_quantizer_tests/layers/max_config.py (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/hints/bits.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/max.rs (1 hunks)
• rust/jstprove_circuits/src/circuit_functions/layers/min.rs (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (2)

rust/jstprove_circuits/src/circuit_functions/hints/bits.rs (2)

python/core/circuits/errors.py (1)

• CircuitError (6-32)

rust/jstprove_circuits/src/circuit_functions/utils/shaping.rs (1)

• input (105-105)

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (3)

rust/jstprove_circuits/src/circuit_functions/utils/onnx_model.rs (3)

• extract_params_and_expected_shape (168-200)
• get_input_name (247-259)
• get_optional_w_or_b (123-132)

rust/jstprove_circuits/src/circuit_functions/utils/tensor_ops.rs (1)

• broadcast_two_arrays (223-238)

rust/jstprove_circuits/src/circuit_functions/gadgets/max_min_clip.rs (3)

• constrained_min (285-344)
• logup_ctx (208-209)
• logup_ctx (330-331)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)

• GitHub Check: Check Formatting and Linting
• GitHub Check: End-to-End Testing (3.10)
• GitHub Check: End-to-End Testing (3.12)
• GitHub Check: End-to-End Testing (3.11)

🔇 Additional comments (4)

python/tests/onnx_quantizer_tests/layers/max_config.py (1)

18-92: MaxConfigProvider and specs look consistent and well‑covered

The Max layer config mirrors the other elementwise layer configs: shapes and names line up, initializers are explicitly cast to float32, and edge/e2e/initializer/broadcast cases (including the empty‑tensor shape) are all exercised. I don’t see further issues here.

rust/jstprove_circuits/src/circuit_functions/layers/min.rs (1)

49-189: MinLayer correctly mirrors MaxLayer using constrained_min and Min‑specific errors

MinLayer cleanly follows the same pattern as MaxLayer: safe input name resolution, initializer/runtime loading, broadcasting, creation of a shared ShiftRangeContext/LogupRangeCheckContext, per‑element constrained_min calls, and finalization plus reshape with InvalidShape diagnostics. LayerKind::Min is used consistently in all error cases, including MissingInput, which keeps diagnostics accurate. I don’t see any issues here.

rust/jstprove_circuits/src/circuit_functions/hints/bits.rs (2)

1-5: LGTM: Clear documentation of unconstrained nature.

The module documentation appropriately emphasizes that these helpers are witness-only and not soundness-critical, which is essential for developers to understand when choosing between constrained and unconstrained operations.

---

105-117: LGTM: Correct bit extraction implementation.

The extraction loop correctly implements the documented algorithm: extract LSB via unconstrained_bit_and, append to results, then shift right. The use of with_capacity pre-allocation is a good efficiency practice, and the resulting little-endian bit order matches the documentation.

[jsgold-1]

Looks awesome 🔥 !

Just the e2e broadcasting tests to ensure we are handling those properly and this LGTM!

[jsgold-1]

LGTM!