JSTprove 1.2.0__py3-none-macosx_11_0_arm64.whl → 1.4.0__py3-none-macosx_11_0_arm64.whl

python/core/model_processing/onnx_quantizer/layers/maxpool.py

@@ -3,9 +3,12 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from typing import ClassVar
+
     import onnx
 
 from python.core.model_processing.onnx_custom_ops.onnx_helpers import (
+    extract_attributes,
     get_attribute_ints,
 )
 from python.core.model_processing.onnx_quantizer.exceptions import InvalidParamError
@@ -21,6 +24,12 @@ class QuantizeMaxpool(QuantizerBase):
     USE_WB = False
     USE_SCALING = False
 
+    DEFAULT_ATTRS: ClassVar = {
+        "dilations": [1],
+        "pads": [0],
+        "strides": [1],
+    }
+
 
 class MaxpoolQuantizer(BaseOpQuantizer, QuantizeMaxpool):
     """
@@ -72,6 +81,35 @@ class MaxpoolQuantizer(BaseOpQuantizer, QuantizeMaxpool):
             InvalidParamError: If any requirement is not met.
         """
         _ = initializer_map
+        attributes = extract_attributes(node)
+        ceil_mode = attributes.get("ceil_mode", None)
+        auto_pad = attributes.get("auto_pad", None)
+        storage_order = attributes.get("storage_order", None)
+
+        if ceil_mode != 0 and ceil_mode is not None:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                "ceil_mode must be 0",
+                "ceil_mode",
+                "0",
+            )
+        if auto_pad != "NOTSET" and auto_pad is not None:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                "auto_pad must be NOTSET",
+                "auto_pad",
+                "NOTSET",
+            )
+        if storage_order != 0 and storage_order is not None:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                "storage_order must be 0",
+                "storage_order",
+                "0",
+            )
         self.check_all_params_exist(node)
         self.check_params_size(node)
         self.check_pool_pads(node)
@@ -85,8 +123,8 @@ class MaxpoolQuantizer(BaseOpQuantizer, QuantizeMaxpool):
         Raises:
             InvalidParamError: If shape requirement is not met.
         """
-        # May need: ["strides", "kernel_shape", "pads", "dilations"]
-        required_attrs = ["strides", "kernel_shape"]
+        required_attrs = ["kernel_shape"]
+
         self.validate_required_attrs(node, required_attrs)
 
         # Check dimension of kernel
@@ -121,11 +159,23 @@ class MaxpoolQuantizer(BaseOpQuantizer, QuantizeMaxpool):
 
     def check_pool_pads(self: MaxpoolQuantizer, node: onnx.NodeProto) -> None:
         kernel_shape = get_attribute_ints(node, "kernel_shape", default=[])
-        pads = get_attribute_ints(node, "pads", default=None)
+        pads_raw = get_attribute_ints(
+            node,
+            "pads",
+            default=self.DEFAULT_ATTRS.get("pads", None),
+        )
+        pads = self.adjust_pads(node, pads_raw)
+
         if pads is None:
             return
         num_dims = len(kernel_shape)
-        if len(pads) != num_dims * 2:
+
+        if len(pads) == 1:
+            pads = pads * 2 * num_dims
+        elif len(pads) == num_dims:
+            # If only beginning pads given, repeat for end pads
+            pads = pads + pads
+        elif len(pads) != num_dims * 2:
             raise InvalidParamError(
                 node.name,
                 node.op_type,
@@ -148,3 +198,28 @@ class MaxpoolQuantizer(BaseOpQuantizer, QuantizeMaxpool):
                     node.op_type,
                     f"pads[{dim + num_dims}]={pad_after} >= kernel[{dim}]={kernel}",
                 )
+
+    def adjust_pads(
+        self: MaxpoolQuantizer,
+        node: onnx.NodeProto,
+        pads_raw: str | int | list[int] | None,
+    ) -> list[int]:
+        if pads_raw is None:
+            pads: list[int] = []
+        elif isinstance(pads_raw, str):
+            # single string, could be "0" or "1 2"
+            pads = [int(x) for x in pads_raw.split()]
+        elif isinstance(pads_raw, int):
+            # single integer
+            pads = [pads_raw]
+        elif isinstance(pads_raw, (list, tuple)):
+            # already a list of numbers (may be strings)
+            pads = [int(x) for x in pads_raw]
+        else:
+            raise InvalidParamError(
+                node.name,
+                node.op_type,
+                f"Cannot parse pads: {pads_raw}",
+            )
+
+        return pads

python/core/model_processing/onnx_quantizer/layers/min.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import onnx
+
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    QuantizerBase,
+    ScaleConfig,
+)
+
+
+class QuantizeMin(QuantizerBase):
+    OP_TYPE = "Min"
+    DOMAIN = ""  # standard ONNX domain
+    USE_WB = True  # let framework wire inputs/outputs normally
+    USE_SCALING = False  # passthrough: no internal scaling
+    SCALE_PLAN: ClassVar = {1: 1}  # elementwise arity plan
+
+
+class MinQuantizer(BaseOpQuantizer, QuantizeMin):
+    """
+    Passthrough quantizer for elementwise Min.
+    We rely on the converter to quantize graph inputs; no extra scaling here.
+    """
+
+    def __init__(
+        self: MinQuantizer,
+        new_initializers: list[onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
+
+    def quantize(
+        self: MinQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        # Delegate to QuantizerBase's generic passthrough implementation.
+        return QuantizeMin.quantize(self, node, graph, scale_config, initializer_map)
+
+    def check_supported(
+        self: MinQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        # Min has no attributes; elementwise, variadic ≥ 1 input per ONNX spec.
+        # We mirror Add/Max broadcasting behavior; no extra checks here.
+        _ = node, initializer_map

python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py

@@ -20,12 +20,15 @@ from python.core.model_processing.onnx_quantizer.layers.base import (
 from python.core.model_processing.onnx_quantizer.layers.batchnorm import (
     BatchnormQuantizer,
 )
+from python.core.model_processing.onnx_quantizer.layers.clip import ClipQuantizer
 from python.core.model_processing.onnx_quantizer.layers.constant import (
     ConstantQuantizer,
 )
 from python.core.model_processing.onnx_quantizer.layers.conv import ConvQuantizer
 from python.core.model_processing.onnx_quantizer.layers.gemm import GemmQuantizer
+from python.core.model_processing.onnx_quantizer.layers.max import MaxQuantizer
 from python.core.model_processing.onnx_quantizer.layers.maxpool import MaxpoolQuantizer
+from python.core.model_processing.onnx_quantizer.layers.min import MinQuantizer
 from python.core.model_processing.onnx_quantizer.layers.mul import MulQuantizer
 from python.core.model_processing.onnx_quantizer.layers.relu import ReluQuantizer
 from python.core.model_processing.onnx_quantizer.layers.sub import SubQuantizer
@@ -74,6 +77,7 @@ class ONNXOpQuantizer:
 
         # Register handlers
         self.register("Add", AddQuantizer(self.new_initializers))
+        self.register("Clip", ClipQuantizer(self.new_initializers))
         self.register("Sub", SubQuantizer(self.new_initializers))
         self.register("Mul", MulQuantizer(self.new_initializers))
         self.register("Conv", ConvQuantizer(self.new_initializers))
@@ -83,6 +87,8 @@ class ONNXOpQuantizer:
         self.register("Constant", ConstantQuantizer())
         self.register("MaxPool", MaxpoolQuantizer())
         self.register("Flatten", PassthroughQuantizer())
+        self.register("Max", MaxQuantizer(self.new_initializers))
+        self.register("Min", MinQuantizer(self.new_initializers))
         self.register("BatchNormalization", BatchnormQuantizer(self.new_initializers))
 
     def register(

python/core/model_templates/circuit_template.py

@@ -1,57 +1,67 @@
-import torch.nn as nn
+from __future__ import annotations
+
+from secrets import randbelow
+
 from python.core.circuits.base import Circuit
-from random import randint
+
 
 class SimpleCircuit(Circuit):
-    '''
-    Note: This template is irrelevant if using the onnx circuit builder. 
-    The template only helps developers if they choose to incorporate other circuit builders into the framework.
+    """
+    Note: This template is irrelevant if using the ONNX circuit builder.
+    The template only helps developers if they choose to incorporate other circuit
+    builders into the framework.
 
-    To begin, we need to specify some basic attributes surrounding the circuit we will be using.
-    required_keys - specify the variables in the input dictionary (and input file).
-    name - name of the rust bin to be run by the circuit.
+    To begin, we need to specify some basic attributes surrounding the circuit we will
+    be using.
 
-    scale_base - specify the base of the scaling applied to each value
-    scale_exponent - the exponent applied to the base to get the scaling factor. Scaling factor will be multiplied by each input
+    - `required_keys`: the variables in the input dictionary (and input file).
+    - `name`: name of the Rust bin to be run by the circuit.
+    - `scale_base`: base of the scaling applied to each value.
+    - `scale_exponent`: exponent applied to the base to get the scaling factor.
+      Scaling factor will be multiplied by each input.
 
-    Other default inputs can be defined below
-    '''
-    def __init__(self, file_name):
+    Other default inputs can be defined below.
+    """
+
+    def __init__(self, file_name: str | None = None) -> None:
         # Initialize the base class
         super().__init__()
-        
+        self.file_name = file_name
+
         # Circuit-specific parameters
         self.required_keys = ["input_a", "input_b", "nonce"]
         self.name = "simple_circuit"  # Use exact name that matches the binary
 
         self.scale_exponent = 1
         self.scale_base = 1
-    
+
         self.input_a = 100
         self.input_b = 200
-        self.nonce = randint(0,10000)
-    
-    '''
-    The following are some important functions used by the model. get inputs should be defined to specify the inputs to the circuit
-    '''
-    def get_inputs(self):
-        '''
-        Specify the inputs to the circuit, based on what was specified in the __init__. Can also have inputs to this function for the inputs.
-        '''
-        return {'input_a': self.input_a, 'input_b': self.input_b, 'nonce': self.nonce}
-    
-    def get_outputs(self, inputs = None):
+        self.nonce = randbelow(10_000)
+
+    def get_inputs(self) -> dict[str, int]:
+        """
+        Specify the inputs to the circuit, based on what was specified
+        in `__init__`.
+        """
+        return {
+            "input_a": self.input_a,
+            "input_b": self.input_b,
+            "nonce": self.nonce,
+        }
+
+    def get_outputs(self, inputs: dict[str, int] | None = None) -> int:
         """
         Compute the output of the circuit.
-        This is overwritten from the base class to ensure computation happens only once.
+
+        This is overwritten from the base class to ensure computation happens
+        only once.
         """
-        if inputs == None:
-            inputs = {'input_a': self.input_a, 'input_b': self.input_b, 'nonce': self.nonce}
-        print(f"Performing addition operation: {inputs['input_a']} + {inputs['input_b']}")
-        return inputs['input_a'] + inputs['input_b']
-    
-    # def format_inputs(self, inputs):
-    #     return {"input": inputs.long().tolist()}
-    
-    # def format_outputs(self, outputs):
-    #     return {"output": outputs.long().tolist()}
+        if inputs is None:
+            inputs = {
+                "input_a": self.input_a,
+                "input_b": self.input_b,
+                "nonce": self.nonce,
+            }
+
+        return inputs["input_a"] + inputs["input_b"]

python/core/utils/errors.py

@@ -30,7 +30,7 @@ class FileCacheError(CircuitExecutionError):
 class ProofBackendError(CircuitExecutionError):
     """Raised when a Cargo command fails."""
 
-    def __init__(  # noqa: PLR0913
+    def __init__(
         self: ProofBackendError,
         message: str,
         command: list[str] | None = None,

python/core/utils/scratch_tests.py

@@ -1,66 +1,72 @@
+from __future__ import annotations
+
 import onnx
-from onnx import TensorProto, helper, shape_inference
-from onnx import numpy_helper
-from onnx import load, save
+from onnx import TensorProto, helper, load, shape_inference
 from onnx.utils import extract_model
 
-def prune_model(model_path, output_names, save_path):
+
+def prune_model(
+    model_path: str,
+    output_names: list[str],
+    save_path: str,
+) -> None:
+    """Extract a sub-model with the same inputs and new outputs."""
     model = load(model_path)
 
-    # Provide model input names and the new desired output names
+    # Provide model input names and the new desired output names.
     input_names = [i.name for i in model.graph.input]
 
     extract_model(
         input_path=model_path,
         output_path=save_path,
         input_names=input_names,
-        output_names=output_names
+        output_names=output_names,
     )
 
-    print(f"Pruned model saved to {save_path}")
+    print(f"Pruned model saved to {save_path}")  # noqa: T201
 
 
-def cut_model(model_path, output_names, save_path):
+def cut_model(
+    model_path: str,
+    output_names: list[str],
+    save_path: str,
+) -> None:
+    """Replace the graph outputs with the tensors named in `output_names`."""
     model = onnx.load(model_path)
     model = shape_inference.infer_shapes(model)
 
     graph = model.graph
 
-    # Remove all current outputs one by one (cannot use .clear() or assignment)
-    while len(graph.output) > 0:
+    # Remove all current outputs one by one (cannot use .clear() or assignment).
+    while graph.output:
         graph.output.pop()
 
-    # Add new outputs
+    # Add new outputs.
     for name in output_names:
-        # Look in value_info, input, or output
+        # Look in value_info, input, or output.
         candidates = list(graph.value_info) + list(graph.input) + list(graph.output)
         value_info = next((vi for vi in candidates if vi.name == name), None)
         if value_info is None:
-            raise ValueError(f"Tensor {name} not found in model graph.")
+            msg = f"Tensor {name} not found in model graph."
+            raise ValueError(msg)
 
         elem_type = value_info.type.tensor_type.elem_type
         shape = [dim.dim_value for dim in value_info.type.tensor_type.shape.dim]
         new_output = helper.make_tensor_value_info(name, elem_type, shape)
         graph.output.append(new_output)
+
     for output in graph.output:
-        print(output)
+        print(output)  # noqa: T201
         if output.name == "/conv1/Conv_output_0":
             output.type.tensor_type.elem_type = TensorProto.INT64
 
     onnx.save(model, save_path)
-    print(f"Saved cut model with outputs {output_names} to {save_path}")
+    print(f"Saved cut model with outputs {output_names} to {save_path}")  # noqa: T201
 
 
 if __name__ == "__main__":
-    # /conv1/Conv_output_0
-    # prune_model(
-    #     model_path="models_onnx/doom.onnx",
-    #     output_names=["/Relu_2_output_0"],  # replace with your intermediate tensor
-    #     save_path= "models_onnx/test_doom_cut.onnx"
-    # )
-    # cut_model("models_onnx/doom.onnx",["/Relu_2_output_0"], "test_doom_after_conv.onnx")
     prune_model(
         model_path="models_onnx/doom.onnx",
         output_names=["/Relu_3_output_0"],  # replace with your intermediate tensor
-        save_path= "models_onnx/test_doom_cut.onnx"
+        save_path="models_onnx/test_doom_cut.onnx",
     )

python/tests/circuit_e2e_tests/circuit_model_developer_test.py

@@ -2,7 +2,11 @@ from __future__ import annotations
 
 import json
 from pathlib import Path
-from typing import Any, Generator
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
 
 import pytest
 import torch
@@ -35,7 +39,7 @@ OUTPUTTWICE = 2
 OUTPUTTHREETIMES = 3
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_circuit_compiles(model_fixture: dict[str, Any]) -> None:
     # Here you could just check that circuit file exists
     circuit_compile_results[model_fixture["model"]] = False
@@ -43,7 +47,7 @@ def test_circuit_compiles(model_fixture: dict[str, Any]) -> None:
     circuit_compile_results[model_fixture["model"]] = True
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_dev(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -89,7 +93,7 @@ def test_witness_dev(
     witness_generated_results[model_fixture["model"]] = True
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_wrong_outputs_dev(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -145,7 +149,7 @@ def test_witness_wrong_outputs_dev(
         ), f"Expected '{output}' in stdout, but it was not found."
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_prove_verify_true_inputs_dev(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -239,7 +243,7 @@ def test_witness_prove_verify_true_inputs_dev(
     ), "Expected 'Verified' in stdout three times, but it was not found."
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_prove_verify_true_inputs_dev_expander_call(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -339,7 +343,7 @@ def test_witness_prove_verify_true_inputs_dev_expander_call(
     assert stdout.count("proving") == 1, "Expected 'proving' but it was not found."
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_read_after_write_json(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -421,7 +425,7 @@ def test_witness_read_after_write_json(
     ), "Input JSON read is not identical to what was written"
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_fresh_compile_dev(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -467,7 +471,7 @@ def test_witness_fresh_compile_dev(
 
 
 # Use once fixed input shape read in rust
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_incorrect_input_shape(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -555,7 +559,7 @@ def test_witness_incorrect_input_shape(
         ), f"Did not expect '{output}' in stdout, but it was found."
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_unscaled(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -666,7 +670,7 @@ def test_witness_unscaled(
     )
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_unscaled_and_incorrect_shape_input(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -790,7 +794,7 @@ def test_witness_unscaled_and_incorrect_shape_input(
     )
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_unscaled_and_incorrect_and_bad_named_input(  # noqa: PLR0915
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -937,7 +941,7 @@ def test_witness_unscaled_and_incorrect_and_bad_named_input(  # noqa: PLR0915
     )
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_wrong_name(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -1063,7 +1067,7 @@ def add_to_first_scalar(data: list, delta: float = 0.1) -> bool:
     return False
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_witness_prove_verify_false_inputs_dev(
     model_fixture: dict[str, Any],
     capsys: Generator[pytest.CaptureFixture[str], None, None],

python/tests/circuit_e2e_tests/helper_fns_for_tests.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Generator, Mapping, Sequence
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Generator, Mapping, Sequence, TypeAlias, Union
+from typing import TYPE_CHECKING, Any, TypeAlias
 
 import numpy as np
 import pytest
@@ -64,7 +65,7 @@ def model_fixture(
     }
 
 
-@pytest.fixture()
+@pytest.fixture
 def temp_witness_file(tmp_path: str) -> Generator[Path, None, None]:
     witness_path = tmp_path / "temp_witness.txt"
     # Give it to the test
@@ -75,7 +76,7 @@ def temp_witness_file(tmp_path: str) -> Generator[Path, None, None]:
         witness_path.unlink()
 
 
-@pytest.fixture()
+@pytest.fixture
 def temp_input_file(tmp_path: str) -> Generator[Path, None, None]:
     input_path = tmp_path / "temp_input.txt"
     # Give it to the test
@@ -86,7 +87,7 @@ def temp_input_file(tmp_path: str) -> Generator[Path, None, None]:
         input_path.unlink()
 
 
-@pytest.fixture()
+@pytest.fixture
 def temp_output_file(tmp_path: str) -> Generator[Path, None, None]:
     output_path = tmp_path / "temp_output.txt"
     # Give it to the test
@@ -97,7 +98,7 @@ def temp_output_file(tmp_path: str) -> Generator[Path, None, None]:
         output_path.unlink()
 
 
-@pytest.fixture()
+@pytest.fixture
 def temp_proof_file(tmp_path: str) -> Generator[Path, None, None]:
     output_path = tmp_path / "temp_proof.txt"
     # Give it to the test
@@ -108,13 +109,10 @@ def temp_proof_file(tmp_path: str) -> Generator[Path, None, None]:
         output_path.unlink()
 
 
-ScalarOrTensor: TypeAlias = Union[int, float, torch.Tensor]
-NestedArray: TypeAlias = Union[
-    ScalarOrTensor,
-    list["NestedArray"],
-    tuple["NestedArray"],
-    np.ndarray,
-]
+ScalarOrTensor: TypeAlias = int | float | torch.Tensor
+NestedArray: TypeAlias = (
+    ScalarOrTensor | list["NestedArray"] | tuple["NestedArray"] | np.ndarray
+)
 
 
 def add_1_to_first_element(x: NestedArray) -> NestedArray:
@@ -137,7 +135,7 @@ def add_1_to_first_element(x: NestedArray) -> NestedArray:
 circuit_compile_results = {}
 witness_generated_results = {}
 
-Nested: TypeAlias = Union[float, Mapping[str, "Nested"], Sequence["Nested"]]
+Nested: TypeAlias = float | Mapping[str, "Nested"] | Sequence["Nested"]
 
 
 def contains_float(obj: Nested) -> bool: