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

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:

python/tests/circuit_parent_classes/test_ort_custom_layers.py

@@ -1,115 +1,97 @@
-import pytest
+from pathlib import Path
+
 import numpy as np
-import torch
 import onnx
-
-from onnx import TensorProto, shape_inference, helper, numpy_helper
+import pytest
+import torch
+from onnx import TensorProto, helper, shape_inference
 
 from python.core.model_processing.converters.onnx_converter import ONNXConverter
-from python.core.model_processing.onnx_custom_ops.onnx_helpers import extract_shape_dict
-from python.core.model_processing.onnx_quantizer.onnx_op_quantizer import ONNXOpQuantizer
-
-from onnxruntime import InferenceSession, SessionOptions
-from onnxruntime_extensions import get_library_path, OrtPyFunction
-from python.core.model_processing.onnx_custom_ops import conv
-
-from python.core.model_processing.onnx_custom_ops.conv import int64_conv
-from python.core.model_processing.onnx_custom_ops.gemm import int64_gemm7
 
 
 @pytest.fixture
-def tiny_conv_model_path(tmp_path):
+def tiny_conv_model_path(tmp_path: Path) -> Path:
     # Create input and output tensor info
-    input_tensor = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1, 1, 4, 4])
-    output_tensor = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1, 1, 2, 2])
+    input_tensor = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 4, 4])
+    output_tensor = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 1, 2, 2])
 
     # Kernel weights (3x3 ones)
-    W_init = helper.make_tensor(
-        name='W',
+    w_init = helper.make_tensor(
+        name="W",
         data_type=TensorProto.FLOAT,
         dims=[1, 1, 3, 3],
-        vals=np.ones((1 * 1 * 3 * 3), dtype=np.float32).tolist()
+        vals=np.ones((1 * 1 * 3 * 3), dtype=np.float32).tolist(),
     )
-    Z_init = helper.make_tensor(
-        name='Z',
+    z_init = helper.make_tensor(
+        name="Z",
         data_type=TensorProto.FLOAT,
         dims=[1],
-        vals=np.ones(( 1), dtype=np.float32).tolist()
+        vals=np.ones((1), dtype=np.float32).tolist(),
     )
 
     # Conv node with no padding, stride 1
     conv_node = helper.make_node(
-        'Conv',
-        inputs=['X', 'W', 'Z'],
-        outputs=['Y'],
+        "Conv",
+        inputs=["X", "W", "Z"],
+        outputs=["Y"],
         kernel_shape=[3, 3],
         pads=[0, 0, 0, 0],
         strides=[1, 1],
-        dilations = [1,1]
+        dilations=[1, 1],
     )
 
     # Build graph and model
     graph = helper.make_graph(
         nodes=[conv_node],
-        name='TinyConvGraph',
+        name="TinyConvGraph",
         inputs=[input_tensor],
         outputs=[output_tensor],
-        initializer=[W_init, Z_init]
+        initializer=[w_init, z_init],
     )
 
-    model = helper.make_model(graph, producer_name='tiny-conv-example')
+    model = helper.make_model(graph, producer_name="tiny-conv-example")
 
     # Save to a temporary file
     model_path = tmp_path / "tiny_conv.onnx"
     onnx.save(model, str(model_path))
 
-    return str(model_path)
+    return model_path
+
 
 @pytest.mark.integration
-def test_tiny_conv(tiny_conv_model_path):
+def test_tiny_conv(tiny_conv_model_path: Path, tmp_path: Path) -> None:
     path = tiny_conv_model_path
 
     converter = ONNXConverter()
 
-    X_input = np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4)
-    id_count = 0
+    # Load and validate original model
     model = onnx.load(path)
-    # Fix, can remove this next line 
-    onnx.checker.check_model(model)
-
-    # Check the model and print Y"s shape information
     onnx.checker.check_model(model)
-    print(f"Before shape inference, the shape info of Y is:\n{model.graph.value_info}")
 
-    # Apply shape inference on the model
+    # Apply shape inference and validate
     inferred_model = shape_inference.infer_shapes(model)
-
-    # Check the model and print Y"s shape information
     onnx.checker.check_model(inferred_model)
-    # print(f"After shape inference, the shape info of Y is:\n{inferred_model.graph.value_info}")
-    
-
-    domain_to_version = {opset.domain: opset.version for opset in model.opset_import}
-    
-    inferred_model = shape_inference.infer_shapes(model)
-    output_name_to_shape = extract_shape_dict(inferred_model)
-    id_count = 0
 
+    # Quantize and add custom domain
     new_model = converter.quantize_model(model, 2, 21)
     custom_domain = onnx.helper.make_operatorsetid(domain="ai.onnx.contrib", version=1)
     new_model.opset_import.append(custom_domain)
     onnx.checker.check_model(new_model)
 
-    with open("model.onnx", "wb") as f:
+    # Save quantized model
+    out_path = tmp_path / "model_quant.onnx"
+    with out_path.open("wb") as f:
         f.write(new_model.SerializeToString())
 
-    model = onnx.load("model.onnx")
-    onnx.checker.check_model(model)  # This throws a descriptive error
+    # Reload quantized model to ensure it is valid
+    model_quant = onnx.load(str(out_path))
+    onnx.checker.check_model(model_quant)
 
+    # 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_quant = converter.run_model_onnx_runtime("model.onnx", inputs)
     true = torch.tensor(np.array(outputs_true), dtype=torch.float32)
     quant = torch.tensor(np.array(outputs_quant), dtype=torch.float32) / (2**21)
 

python/tests/onnx_quantizer_tests/layers/base.py

@@ -113,9 +113,7 @@ class LayerTestConfig:
         # respect that; otherwise use original valid_inputs.
         inputs = test_spec.input_overrides or self.valid_inputs
 
-        # Prepare attributes
-        attrs = {**self.valid_attributes, **test_spec.attr_overrides}
-        # Remove omitted attributes if specified
+        # Prepare attributes and remove omitted attributes if specified
         attrs = {**self.valid_attributes, **test_spec.attr_overrides}
         for key in getattr(test_spec, "omit_attrs", []):
             attrs.pop(key, None)