JSTprove 1.0.0__py3-none-macosx_11_0_arm64.whl → 1.1.0__py3-none-macosx_11_0_arm64.whl

python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
-from typing import Callable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
 
 import onnx
 
@@ -9,6 +12,7 @@ from python.core.model_processing.onnx_quantizer.exceptions import (
     MissingHandlerError,
     UnsupportedOpError,
 )
+from python.core.model_processing.onnx_quantizer.layers.add import AddQuantizer
 from python.core.model_processing.onnx_quantizer.layers.base import (
     PassthroughQuantizer,
     ScaleConfig,
@@ -64,6 +68,7 @@ class ONNXOpQuantizer:
         self.new_initializers = []
 
         # Register handlers
+        self.register("Add", AddQuantizer(self.new_initializers))
         self.register("Conv", ConvQuantizer(self.new_initializers))
         self.register("Relu", ReluQuantizer())
         self.register("Reshape", PassthroughQuantizer())

python/core/utils/general_layer_functions.py

@@ -83,7 +83,9 @@ class GeneralLayerFunctions:
             tensor = torch.mul(tensor, self.scale_base**self.scale_exponent)
 
         tensor = tensor.long()
+        return self.reshape_inputs(tensor)
 
+    def reshape_inputs(self: GeneralLayerFunctions, tensor: torch.Tensor) -> None:
         if hasattr(self, "input_shape"):
             shape = self.input_shape
             if hasattr(self, "adjust_shape") and callable(
@@ -94,6 +96,7 @@ class GeneralLayerFunctions:
                 tensor = tensor.reshape(shape)
             except RuntimeError as e:
                 raise ShapeMismatchError(shape, list(tensor.shape)) from e
+
         return tensor
 
     def get_inputs(
@@ -154,21 +157,19 @@ class GeneralLayerFunctions:
                 # If unknown dim in batch spot, assume batch size of 1
                 first_key = next(iter(keys))
                 input_shape = self.input_shape[first_key]
-                input_shape[0] = 1 if input_shape[0] < 1 else input_shape[0]
+                input_shape[0] = max(input_shape[0], 1)
                 return self.get_rand_inputs(input_shape)
             inputs = {}
             for key in keys:
                 # If unknown dim in batch spot, assume batch size of 1
-                input_shape = self.input_shape[keys[key]]
-                if not isinstance(input_shape, list) and not isinstance(
-                    input_shape,
-                    tuple,
-                ):
+                input_shape = self.input_shape[key]
+                if not isinstance(input_shape, (list, tuple)):
                     msg = f"Invalid input shape for key '{key}': {input_shape}"
                     raise CircuitUtilsError(msg)
-                input_shape[0] = 1 if input_shape[0] < 1 else input_shape[0]
+                input_shape[0] = max(input_shape[0], 1)
                 inputs[key] = self.get_rand_inputs(input_shape)
             return inputs
+
         if not (hasattr(self, "scale_base") and hasattr(self, "scale_exponent")):
             attr_name = "scale_base/scale_exponent"
             context = "needed for scaling random inputs"
@@ -195,8 +196,10 @@ class GeneralLayerFunctions:
             torch.Tensor: A tensor of random values in [-1, 1).
         """
         if not isinstance(input_shape, (list, tuple)):
-            msg = f"Invalid input_shape type: {type(input_shape)}."
-            " Expected list or tuple of ints."
+            msg = (
+                f"Invalid input_shape type: {type(input_shape)}."
+                " Expected list or tuple of ints."
+            )
             raise CircuitUtilsError(msg)
         if not all(isinstance(x, int) and x > 0 for x in input_shape):
             raise ShapeMismatchError(
@@ -240,9 +243,11 @@ class GeneralLayerFunctions:
             try:
                 rescaled = torch.div(outputs, self.scale_base**self.scale_exponent)
             except Exception as e:
-                msg = "Failed to rescale outputs using scale_base="
-                f"{getattr(self, 'scale_base', None)} "
-                f"and scale_exponent={getattr(self, 'scale_exponent', None)}: {e}"
+                msg = (
+                    "Failed to rescale outputs using scale_base="
+                    f"{getattr(self, 'scale_base', None)} "
+                    f"and scale_exponent={getattr(self, 'scale_exponent', None)}: {e}"
+                )
                 raise CircuitUtilsError(msg) from e
             return {
                 "output": outputs.long().tolist(),

python/core/utils/model_registry.py

@@ -147,7 +147,8 @@ def get_models_to_test(
 
     Args:
         selected_models (list[str], optional):
-            A list of model names to include. Defaults to None.
+            A list of model names to include. If None, returns empty list.
+            If provided but no matches found, returns empty list. Defaults to None.
         source_filter (str, optional):
             Restrict models to a specific source (e.g., "onnx", "class").
             Defaults to None.
@@ -155,10 +156,12 @@ def get_models_to_test(
     Returns:
         list[ModelEntry]: A filtered list of model entries.
     """
+    if selected_models is None:
+        return []
+
     models = MODELS_TO_TEST
 
-    if selected_models is not None:
-        models = [m for m in models if m.name in selected_models]
+    models = [m for m in models if m.name in selected_models]
 
     if source_filter is not None:
         models = [m for m in models if m.source == source_filter]

python/tests/circuit_e2e_tests/other_e2e_test.py

@@ -1,13 +1,15 @@
+# ruff: noqa: S603
 import json
 import subprocess
 import sys
+from collections.abc import Generator
 from pathlib import Path
-from typing import Generator
 
 import numpy as np
 import onnx
 import pytest
-from onnx import helper, numpy_helper
+import torch
+from onnx import TensorProto, helper, numpy_helper
 
 
 def create_simple_gemm_onnx_model(
@@ -19,14 +21,14 @@ def create_simple_gemm_onnx_model(
     # Define input
     input_tensor = helper.make_tensor_value_info(
         "input",
-        onnx.TensorProto.FLOAT,
+        TensorProto.FLOAT,
         [1, input_size],
     )
 
     # Define output
     output_tensor = helper.make_tensor_value_info(
         "output",
-        onnx.TensorProto.FLOAT,
+        TensorProto.FLOAT,
         [1, output_size],
     )
 
@@ -67,7 +69,7 @@ def create_simple_gemm_onnx_model(
     onnx.save(model, str(model_path))
 
 
-@pytest.mark.e2e()
+@pytest.mark.e2e
 def test_parallel_compile_and_witness_two_simple_models(  # noqa: PLR0915
     tmp_path: str,
     capsys: Generator[pytest.CaptureFixture[str], None, None],
@@ -115,7 +117,7 @@ def test_parallel_compile_and_witness_two_simple_models(  # noqa: PLR0915
 
     # Run compile commands
     result1 = subprocess.run(
-        compile_cmd1,  # noqa: S603
+        compile_cmd1,
         capture_output=True,
         text=True,
         check=False,
@@ -123,7 +125,7 @@ def test_parallel_compile_and_witness_two_simple_models(  # noqa: PLR0915
     assert result1.returncode == 0, f"Compile failed for model1: {result1.stderr}"
 
     result2 = subprocess.run(
-        compile_cmd2,  # noqa: S603
+        compile_cmd2,
         capture_output=True,
         text=True,
         check=False,
@@ -179,8 +181,8 @@ def test_parallel_compile_and_witness_two_simple_models(  # noqa: PLR0915
     ]
 
     # Start both processes
-    proc1 = subprocess.Popen(witness_cmd1)  # noqa: S603
-    proc2 = subprocess.Popen(witness_cmd2)  # noqa: S603
+    proc1 = subprocess.Popen(witness_cmd1)
+    proc2 = subprocess.Popen(witness_cmd2)
 
     # Wait for both to complete
     proc1.wait()
@@ -215,3 +217,194 @@ def test_parallel_compile_and_witness_two_simple_models(  # noqa: PLR0915
         len(output2["output"]) == model2_output_size
     ), f"Output2 should have {model2_output_size} elements,"
     f" got {len(output2['output'])}"
+
+
+def create_multi_input_multi_output_model(model_path: Path) -> None:
+    """Create a simple ONNX model with two inputs and two outputs."""
+    # Define inputs
+    x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 4, 4])
+    w = helper.make_tensor_value_info("W", TensorProto.FLOAT, [1, 1, 4, 4])
+
+    # Define outputs
+    y1 = helper.make_tensor_value_info("sum", TensorProto.FLOAT, [1, 1, 4, 4])
+    y2 = helper.make_tensor_value_info("pooled", TensorProto.FLOAT, [1, 1, 2, 2])
+
+    # Node 1: Add
+    add_node = helper.make_node("Add", inputs=["X", "W"], outputs=["sum"])
+
+    # Node 2: MaxPool
+    pool_node = helper.make_node(
+        "MaxPool",
+        inputs=["sum"],
+        outputs=["pooled"],
+        kernel_shape=[2, 2],
+        strides=[2, 2],
+        dilations=[1, 1],
+        pads=[0, 0, 0, 0],
+        ceil_mode=0,
+    )
+
+    # Build the graph
+    graph_def = helper.make_graph(
+        [add_node, pool_node],
+        "TwoOutputGraph",
+        [x, w],
+        [y1, y2],
+    )
+
+    model_def = helper.make_model(graph_def, producer_name="pytest-multi-output-model")
+    onnx.save(model_def, model_path)
+
+
+@pytest.mark.e2e
+def test_multi_input_multi_output_model_e2e(tmp_path: Path) -> None:
+    """
+    E2E test: compile, witness, and verify outputs
+    for a multi-input/multi-output ONNX model.
+    """
+    model_path = tmp_path / "multi_output_no_identity.onnx"
+    circuit_path = tmp_path / "circuit.txt"
+    input_path = tmp_path / "input.json"
+    output_path = tmp_path / "output.json"
+    witness_path = tmp_path / "witness.bin"
+    proof_path = tmp_path / "proof.bin"
+
+    # --- Step 1: Generate model ---
+    create_multi_input_multi_output_model(model_path)
+
+    # --- Step 2: Compile model ---
+    compile_cmd = [
+        sys.executable,
+        "-m",
+        "python.frontend.cli",
+        "compile",
+        "-m",
+        str(model_path),
+        "-c",
+        str(circuit_path),
+    ]
+    result = subprocess.run(compile_cmd, capture_output=True, text=True, check=False)
+    assert (
+        result.returncode == 0
+    ), f"Compile failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
+
+    # --- Step 3: Create input JSON ---
+    # Simple constant tensors (shape [1,1,4,4])
+    x = [
+        [
+            [
+                [1.0, 2.0, 3.0, 4.0],
+                [5.0, 6.0, 7.0, 8.0],
+                [9.0, 10.0, 11.0, 12.0],
+                [13.0, 14.0, 15.0, 16.0],
+            ],
+        ],
+    ]
+    w = [
+        [
+            [
+                [0.1, 0.2, 0.3, 0.4],
+                [0.5, 0.6, 0.7, 0.8],
+                [0.9, 1.0, 1.1, 1.2],
+                [1.3, 1.4, 1.5, 1.6],
+            ],
+        ],
+    ]
+
+    with Path.open(input_path, "w") as f:
+        json.dump({"X": x, "W": w}, f)
+
+    # --- Step 4: Run witness ---
+    witness_cmd = [
+        sys.executable,
+        "-m",
+        "python.frontend.cli",
+        "witness",
+        "-c",
+        str(circuit_path),
+        "-i",
+        str(input_path),
+        "-o",
+        str(output_path),
+        "-w",
+        str(witness_path),
+    ]
+    result = subprocess.run(witness_cmd, capture_output=True, text=True, check=False)
+    assert (
+        result.returncode == 0
+    ), f"Witness failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
+
+    # --- Step 5: Validate output files ---
+    assert output_path.exists(), "Output file not generated"
+    assert witness_path.exists(), "Witness file not generated"
+
+    with Path.open(output_path) as f:
+        outputs = json.load(f)
+
+    output_raw = (
+        (torch.as_tensor(x) * 2**18).long() + (torch.as_tensor(w) * 2**18).long()
+    ).flatten()
+
+    second_outputs = output_raw.clone().reshape([1, 1, 4, 4])
+
+    outputs_2 = torch.max_pool2d(
+        second_outputs,
+        kernel_size=2,
+        stride=2,
+        dilation=1,
+        padding=0,
+    ).flatten()
+
+    output_raw = torch.cat((output_raw, outputs_2))
+
+    assert torch.allclose(
+        torch.as_tensor(outputs["output"]),
+        output_raw,
+        rtol=1e-3,
+        atol=1e-5,
+    ), "Outputs do not match"
+
+    # --- Step 5: Prove ---
+    prove_cmd = [
+        sys.executable,
+        "-m",
+        "python.frontend.cli",
+        "prove",
+        "-c",
+        str(circuit_path),
+        "-w",
+        str(witness_path),
+        "-p",
+        str(proof_path),
+    ]
+    result = subprocess.run(prove_cmd, check=False, capture_output=True, text=True)
+    assert (
+        result.returncode == 0
+    ), f"Prove failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
+
+    # --- Step 6: Verify ---
+    verify_cmd = [
+        sys.executable,
+        "-m",
+        "python.frontend.cli",
+        "verify",
+        "-c",
+        str(circuit_path),
+        "-i",
+        str(input_path),
+        "-o",
+        str(output_path),
+        "-w",
+        str(witness_path),
+        "-p",
+        str(proof_path),
+    ]
+    result = subprocess.run(verify_cmd, check=False, capture_output=True, text=True)
+    assert (
+        result.returncode == 0
+    ), f"Verify failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
+
+    # --- Step 7: Validate output ---
+    assert output_path.exists(), "Output JSON not generated"
+    assert witness_path.exists(), "Witness not generated"
+    assert proof_path.exists(), "Proof not generated"