JSTprove 1.0.0__py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

python/tests/circuit_e2e_tests/helper_fns_for_tests.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Mapping, Sequence, TypeAlias, Union
+
+import numpy as np
+import pytest
+import torch
+
+if TYPE_CHECKING:
+    from python.core.circuits.zk_model_base import ZKModelBase
+from python.core.utils.helper_functions import CircuitExecutionConfig, RunType
+
+GOOD_OUTPUT = ["Witness Generated"]
+BAD_OUTPUT = [
+    "Witness generation failed",
+    "Outputs generated do not match outputs supplied",
+]
+
+NUMPARAMS3 = 3
+NUMPARAMS4 = 4
+
+
+@pytest.fixture(scope="module")
+def model_fixture(
+    request: pytest.FixtureRequest,
+    tmp_path_factory: pytest.TempPathFactory,
+) -> dict[str, Any]:
+    param = request.param
+    name = f"{param.name}"
+    model_class = param.loader
+    args, kwargs = (), {}
+
+    if len(param) == NUMPARAMS3:
+        if isinstance(param[2], dict):
+            kwargs = param[2]
+        else:
+            args = param[2]
+    elif len(param) == NUMPARAMS4:
+        args, kwargs = param[2], param[3]
+
+    temp_dir = tmp_path_factory.mktemp(name)
+    circuit_path = temp_dir / f"{name}_circuit.txt"
+    quantized_path = temp_dir / f"{name}_quantized.pt"
+
+    model = model_class(*args, **kwargs)
+
+    model.base_testing(
+        CircuitExecutionConfig(
+            run_type=RunType.COMPILE_CIRCUIT,
+            dev_mode=True,
+            circuit_path=str(circuit_path),
+            quantized_path=quantized_path,
+        ),
+    )
+
+    return {
+        "name": name,
+        "model_class": model_class,
+        "circuit_path": circuit_path,
+        "temp_dir": temp_dir,
+        "model": model,
+        "quantized_model": quantized_path,
+    }
+
+
+@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
+    yield witness_path
+
+    # After the test is done, remove it
+    if Path.exists(witness_path):
+        witness_path.unlink()
+
+
+@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
+    yield input_path
+
+    # After the test is done, remove it
+    if Path.exists(input_path):
+        input_path.unlink()
+
+
+@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
+    yield output_path
+
+    # After the test is done, remove it
+    if Path.exists(output_path):
+        output_path.unlink()
+
+
+@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
+    yield output_path
+
+    # After the test is done, remove it
+    if Path.exists(output_path):
+        output_path.unlink()
+
+
+ScalarOrTensor: TypeAlias = Union[int, float, torch.Tensor]
+NestedArray: TypeAlias = Union[
+    ScalarOrTensor,
+    list["NestedArray"],
+    tuple["NestedArray"],
+    np.ndarray,
+]
+
+
+def add_1_to_first_element(x: NestedArray) -> NestedArray:
+    """Safely adds 1 to the first element of any scalar/list/tensor."""
+    if isinstance(x, (int, float)):
+        return x + 1
+    if isinstance(x, torch.Tensor):
+        x = x.clone()  # avoid in-place modification
+        x.view(-1)[0] += 1
+        return x
+    if isinstance(x, (list, tuple, np.ndarray)):
+        x = list(x)
+        x[0] = add_1_to_first_element(x[0])
+        return x
+    msg = f"Unsupported type for get_outputs patch: {type(x)}"
+    raise TypeError(msg)
+
+
+# Define models to be tested
+circuit_compile_results = {}
+witness_generated_results = {}
+
+Nested: TypeAlias = Union[float, Mapping[str, "Nested"], Sequence["Nested"]]
+
+
+def contains_float(obj: Nested) -> bool:
+    if isinstance(obj, float):
+        return True
+    if isinstance(obj, dict):
+        return any(contains_float(v) for v in obj.values())
+    if isinstance(obj, list):
+        return any(contains_float(i) for i in obj)
+    return False
+
+
+@pytest.fixture(scope="module")
+def check_model_compiles(model_fixture: dict[str, Any]) -> None:
+    # Default to True; will be set to False if first test fails
+    result = circuit_compile_results.get(model_fixture["model"])
+    if result is False:
+        pytest.skip(
+            f"Skipping because the first test failed for: {model_fixture['model']}",
+        )
+    return result
+
+
+@pytest.fixture(scope="module")
+def check_witness_generated(model_fixture: dict[str, Any]) -> None:
+    # Default to True; will be set to False if first test fails
+    result = witness_generated_results.get(model_fixture["model"])
+    if result is False:
+        pytest.skip(
+            f"Skipping because the first test failed for: {model_fixture['model']}",
+        )
+    return result
+
+
+def assert_very_close(
+    inputs_1: np.array,
+    inputs_2: np.array,
+    model: ZKModelBase,
+) -> None:
+    for i in inputs_1:
+        x = torch.div(
+            torch.as_tensor(inputs_1[i]),
+            model.scale_base**model.scale_exponent,
+        )
+        y = torch.div(
+            torch.as_tensor(inputs_2[i]),
+            model.scale_base**model.scale_exponent,
+        )
+
+        assert torch.isclose(x, y, rtol=1e-8).all()

python/tests/circuit_e2e_tests/other_e2e_test.py

@@ -0,0 +1,217 @@
+import json
+import subprocess
+import sys
+from pathlib import Path
+from typing import Generator
+
+import numpy as np
+import onnx
+import pytest
+from onnx import helper, numpy_helper
+
+
+def create_simple_gemm_onnx_model(
+    input_size: int,
+    output_size: int,
+    model_path: Path,
+) -> None:
+    """Create a simple ONNX model with a single GEMM layer."""
+    # Define input
+    input_tensor = helper.make_tensor_value_info(
+        "input",
+        onnx.TensorProto.FLOAT,
+        [1, input_size],
+    )
+
+    # Define output
+    output_tensor = helper.make_tensor_value_info(
+        "output",
+        onnx.TensorProto.FLOAT,
+        [1, output_size],
+    )
+
+    # Create random number generator
+    rng = np.random.default_rng()
+
+    # Create weight tensor
+    weight = rng.standard_normal((output_size, input_size)).astype(np.float32)
+    weight_tensor = numpy_helper.from_array(weight, name="weight")
+
+    # Create bias tensor
+    bias = rng.standard_normal((output_size,)).astype(np.float32)
+    bias_tensor = numpy_helper.from_array(bias, name="bias")
+
+    # Create GEMM node
+    gemm_node = helper.make_node(
+        "Gemm",
+        inputs=["input", "weight", "bias"],
+        outputs=["output"],
+        alpha=1.0,
+        beta=1.0,
+        transB=1,  # Transpose B (weight)
+    )
+
+    # Create graph
+    graph = helper.make_graph(
+        [gemm_node],
+        "simple_gemm",
+        [input_tensor],
+        [output_tensor],
+        [weight_tensor, bias_tensor],
+    )
+
+    # Create model
+    model = helper.make_model(graph, producer_name="simple_gemm_creator")
+
+    # Save model
+    onnx.save(model, str(model_path))
+
+
+@pytest.mark.e2e()
+def test_parallel_compile_and_witness_two_simple_models(  # noqa: PLR0915
+    tmp_path: str,
+    capsys: Generator[pytest.CaptureFixture[str], None, None],
+) -> None:
+    """Test compiling and running witness
+    for two different simple ONNX models in parallel.
+    """
+    # Create two simple ONNX models with different shapes
+    model1_path = Path(tmp_path) / "simple_gemm1.onnx"
+    model2_path = Path(tmp_path) / "simple_gemm2.onnx"
+    model1_input_size = 4
+    model1_output_size = 2
+
+    model2_input_size = 6
+    model2_output_size = 3
+
+    create_simple_gemm_onnx_model(model1_input_size, model1_output_size, model1_path)
+    create_simple_gemm_onnx_model(model2_input_size, model2_output_size, model2_path)
+
+    # Define paths for artifacts
+    circuit1_path = Path(tmp_path) / "circuit1.txt"
+    circuit2_path = Path(tmp_path) / "circuit2.txt"
+
+    # Compile both models
+    compile_cmd1 = [
+        sys.executable,
+        "-m",
+        "python.frontend.cli",
+        "compile",
+        "-m",
+        str(model1_path),
+        "-c",
+        str(circuit1_path),
+    ]
+    compile_cmd2 = [
+        sys.executable,
+        "-m",
+        "python.frontend.cli",
+        "compile",
+        "-m",
+        str(model2_path),
+        "-c",
+        str(circuit2_path),
+    ]
+
+    # Run compile commands
+    result1 = subprocess.run(
+        compile_cmd1,  # noqa: S603
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert result1.returncode == 0, f"Compile failed for model1: {result1.stderr}"
+
+    result2 = subprocess.run(
+        compile_cmd2,  # noqa: S603
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert result2.returncode == 0, f"Compile failed for model2: {result2.stderr}"
+
+    # Create input files
+    input1_data = {"input": [1.0] * model1_input_size}  # 10 inputs
+    input2_data = {"input": [1.0] * model2_input_size}  # 20 inputs
+
+    input1_path = Path(tmp_path) / "input1.json"
+    input2_path = Path(tmp_path) / "input2.json"
+
+    with Path.open(input1_path, "w") as f:
+        json.dump(input1_data, f)
+    with Path.open(input2_path, "w") as f:
+        json.dump(input2_data, f)
+
+    # Define output and witness paths
+    output1_path = Path(tmp_path) / "output1.json"
+    witness1_path = Path(tmp_path) / "witness1.bin"
+    output2_path = Path(tmp_path) / "output2.json"
+    witness2_path = Path(tmp_path) / "witness2.bin"
+
+    # Run witness commands in parallel
+    witness_cmd1 = [
+        sys.executable,
+        "-m",
+        "python.frontend.cli",
+        "witness",
+        "-c",
+        str(circuit1_path),
+        "-i",
+        str(input1_path),
+        "-o",
+        str(output1_path),
+        "-w",
+        str(witness1_path),
+    ]
+    witness_cmd2 = [
+        sys.executable,
+        "-m",
+        "python.frontend.cli",
+        "witness",
+        "-c",
+        str(circuit2_path),
+        "-i",
+        str(input2_path),
+        "-o",
+        str(output2_path),
+        "-w",
+        str(witness2_path),
+    ]
+
+    # Start both processes
+    proc1 = subprocess.Popen(witness_cmd1)  # noqa: S603
+    proc2 = subprocess.Popen(witness_cmd2)  # noqa: S603
+
+    # Wait for both to complete
+    proc1.wait()
+    proc2.wait()
+
+    # Check return codes
+    assert proc1.returncode == 0, "Witness failed for model1"
+    assert proc2.returncode == 0, "Witness failed for model2"
+
+    # Verify outputs exist
+    assert output1_path.exists(), "Output1 file not generated"
+    assert output2_path.exists(), "Output2 file not generated"
+    assert witness1_path.exists(), "Witness1 file not generated"
+    assert witness2_path.exists(), "Witness2 file not generated"
+
+    # Check output contents (should have the correct shapes)
+    with Path.open(output1_path) as f:
+        output1 = json.load(f)
+    with Path.open(output2_path) as f:
+        output2 = json.load(f)
+
+    # Model1: input 10 -> output 5
+    assert "output" in output1, "Output1 missing 'output' key"
+    assert (
+        len(output1["output"]) == model1_output_size
+    ), f"Output1 should have {model1_output_size} elements,"
+    f" got {len(output1['output'])}"
+
+    # Model2: input 20 -> output 8
+    assert "output" in output2, "Output2 missing 'output' key"
+    assert (
+        len(output2["output"]) == model2_output_size
+    ), f"Output2 should have {model2_output_size} elements,"
+    f" got {len(output2['output'])}"