JSTprove 1.0.0__py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

python/frontend/commands/bench/sweep.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import argparse
+
+from python.core.utils.helper_functions import run_subprocess
+from python.frontend.commands.args import ArgSpec
+from python.frontend.commands.base import BaseCommand
+
+SWEEP_CHOICES: tuple[str, ...] = ("depth", "breadth", "lenet")
+
+DEPTH_MIN = ArgSpec(
+    name="depth_min",
+    flag="--depth-min",
+    help_text="(depth) minimum conv depth",
+    arg_type=int,
+)
+
+DEPTH_MAX = ArgSpec(
+    name="depth_max",
+    flag="--depth-max",
+    help_text="(depth) maximum conv depth",
+    arg_type=int,
+)
+
+INPUT_HW = ArgSpec(
+    name="input_hw",
+    flag="--input-hw",
+    help_text="(depth) input H=W (e.g., 56)",
+    arg_type=int,
+)
+
+ARCH_DEPTH = ArgSpec(
+    name="arch_depth",
+    flag="--arch-depth",
+    help_text="(breadth) conv blocks at fixed topology",
+    arg_type=int,
+)
+
+INPUT_HW_LIST = ArgSpec(
+    name="input_hw_list",
+    flag="--input-hw-list",
+    help_text="(breadth) sizes like '28,56,84,112' or '32:160:32'",
+)
+
+ITERATIONS = ArgSpec(
+    name="iterations",
+    flag="--iterations",
+    help_text="E2E loops per model (default 3 in sweeps)",
+    arg_type=int,
+)
+
+RESULTS = ArgSpec(
+    name="results",
+    flag="--results",
+    help_text="Path to JSONL results (e.g., benchmarking/depth_sweep.jsonl)",
+)
+
+ONNX_DIR = ArgSpec(
+    name="onnx_dir",
+    flag="--onnx-dir",
+    help_text="Override ONNX output dir (else chosen by sweep)",
+)
+
+INPUTS_DIR = ArgSpec(
+    name="inputs_dir",
+    flag="--inputs-dir",
+    help_text="Override inputs output dir (else chosen by sweep)",
+)
+
+POOL_CAP = ArgSpec(
+    name="pool_cap",
+    flag="--pool-cap",
+    help_text="Max pool blocks at start (Lenet-like: 2)",
+    arg_type=int,
+)
+
+STOP_AT_HW = ArgSpec(
+    name="stop_at_hw",
+    flag="--stop-at-hw",
+    help_text="Allow pooling while H >= this",
+    arg_type=int,
+)
+
+CONV_OUT_CH = ArgSpec(
+    name="conv_out_ch",
+    flag="--conv-out-ch",
+    help_text="Conv output channels",
+    arg_type=int,
+)
+
+FC_HIDDEN = ArgSpec(
+    name="fc_hidden",
+    flag="--fc-hidden",
+    help_text="Fully-connected hidden size",
+    arg_type=int,
+)
+
+N_ACTIONS = ArgSpec(
+    name="n_actions",
+    flag="--n-actions",
+    help_text="Classifier outputs (classes)",
+    arg_type=int,
+)
+
+TAG = ArgSpec(
+    name="tag",
+    flag="--tag",
+    help_text="Optional tag suffix for filenames",
+)
+
+MODE = ArgSpec(
+    name="mode",
+    flag="mode",
+    help_text="Sweep type: depth, breadth, or lenet.",
+    extra_kwargs={"choices": list(SWEEP_CHOICES)},
+)
+
+
+class SweepCommand(BaseCommand):
+    """Generate ONNX models and benchmark parameter sweeps (depth/breadth)."""
+
+    name: ClassVar[str] = "sweep"
+    aliases: ClassVar[list[str]] = []
+    help: ClassVar[str] = (
+        "Generate ONNX models and benchmark parameter sweeps (depth/breadth)."
+    )
+
+    SCRIPT_GEN_AND_BENCH: ClassVar[str] = "python.scripts.gen_and_bench"
+
+    DEPTH_DEFAULTS: ClassVar[tuple[tuple[ArgSpec, str], ...]] = (
+        (DEPTH_MIN, "1"),
+        (DEPTH_MAX, "16"),
+        (ITERATIONS, "3"),
+        (RESULTS, "benchmarking/depth_sweep.jsonl"),
+    )
+
+    BREADTH_DEFAULTS: ClassVar[tuple[tuple[ArgSpec, str], ...]] = (
+        (ARCH_DEPTH, "5"),
+        (INPUT_HW_LIST, "28,56,84,112"),
+        (ITERATIONS, "3"),
+        (RESULTS, "benchmarking/breadth_sweep.jsonl"),
+        (POOL_CAP, "2"),
+        (CONV_OUT_CH, "16"),
+        (FC_HIDDEN, "256"),
+    )
+
+    SWEEP_CONFIGS: ClassVar[dict[str, tuple[tuple[ArgSpec, str], ...]]] = {
+        SWEEP_CHOICES[0]: DEPTH_DEFAULTS,
+        SWEEP_CHOICES[1]: BREADTH_DEFAULTS,
+    }
+
+    SWEEP_ARGS: ClassVar[tuple[ArgSpec, ...]] = (
+        DEPTH_MIN,
+        DEPTH_MAX,
+        INPUT_HW,
+        ARCH_DEPTH,
+        INPUT_HW_LIST,
+        ITERATIONS,
+        RESULTS,
+        ONNX_DIR,
+        INPUTS_DIR,
+        POOL_CAP,
+        STOP_AT_HW,
+        CONV_OUT_CH,
+        FC_HIDDEN,
+        N_ACTIONS,
+        TAG,
+    )
+
+    OPTIONAL_SWEEP_ARGS: ClassVar[tuple[ArgSpec, ...]] = (ONNX_DIR, INPUTS_DIR, TAG)
+
+    @classmethod
+    def configure_parser(
+        cls: type[SweepCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        MODE.add_to_parser(parser)
+
+        for arg_spec in cls.SWEEP_ARGS:
+            arg_spec.add_to_parser(parser)
+
+    @classmethod
+    def run(cls: type[SweepCommand], args: argparse.Namespace) -> None:
+        if not args.mode:
+            msg = "Specify sweep type: depth, breadth, or lenet"
+            raise ValueError(msg)
+
+        cls._run_sweep_benchmark(args, args.mode)
+
+    @classmethod
+    def _run_sweep_benchmark(
+        cls: type[SweepCommand],
+        args: argparse.Namespace,
+        sweep: str,
+    ) -> None:
+        provided_knobs = [getattr(args, spec.name, None) for spec in cls.SWEEP_ARGS]
+        simple = all(v is None for v in provided_knobs)
+
+        cmd = [sys.executable, "-m", cls.SCRIPT_GEN_AND_BENCH, "--sweep", sweep]
+
+        if simple:
+            defaults = cls.SWEEP_CONFIGS.get(sweep, ())
+            BaseCommand.append_args_from_specs(cmd, defaults)
+            BaseCommand.append_args_from_namespace(cmd, args, cls.OPTIONAL_SWEEP_ARGS)
+        else:
+            BaseCommand.append_args_from_namespace(cmd, args, cls.SWEEP_ARGS)
+
+        run_subprocess(cmd)

python/frontend/commands/compile.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import argparse
+
+from python.core.circuits.errors import CircuitRunError
+from python.core.utils.helper_functions import CircuitExecutionConfig, RunType
+from python.frontend.commands.args import CIRCUIT_PATH, MODEL_PATH
+from python.frontend.commands.base import BaseCommand
+
+
+class CompileCommand(BaseCommand):
+    """Compile an ONNX model to a circuit."""
+
+    name: ClassVar[str] = "compile"
+    aliases: ClassVar[list[str]] = ["comp"]
+    help: ClassVar[str] = (
+        "Compile a circuit (writes circuit + quantized model + weights)."
+    )
+
+    @classmethod
+    def configure_parser(
+        cls: type[CompileCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        MODEL_PATH.add_to_parser(parser)
+        CIRCUIT_PATH.add_to_parser(
+            parser,
+            "Output path for the compiled circuit (e.g., circuit.txt).",
+        )
+
+    @classmethod
+    @BaseCommand.validate_required(MODEL_PATH, CIRCUIT_PATH)
+    @BaseCommand.validate_paths(MODEL_PATH)
+    @BaseCommand.validate_parent_paths(CIRCUIT_PATH)
+    def run(cls: type[CompileCommand], args: argparse.Namespace) -> None:
+        model_name_hint = Path(args.model_path).stem
+        circuit = cls._build_circuit(model_name_hint)
+
+        circuit.model_file_name = args.model_path
+        circuit.onnx_path = args.model_path
+        circuit.model_path = args.model_path
+
+        try:
+            circuit.base_testing(
+                CircuitExecutionConfig(
+                    run_type=RunType.COMPILE_CIRCUIT,
+                    circuit_path=args.circuit_path,
+                    dev_mode=False,
+                ),
+            )
+        except CircuitRunError as e:
+            raise RuntimeError(e) from e
+
+        print(f"[compile] done → circuit={args.circuit_path}")  # noqa: T201

python/frontend/commands/constants.py

@@ -0,0 +1,5 @@
+"""Constants for CLI commands."""
+
+# Circuit defaults
+DEFAULT_CIRCUIT_MODULE = "python.core.circuit_models.generic_onnx"
+DEFAULT_CIRCUIT_CLASS = "GenericModelONNX"

python/frontend/commands/model_check.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import argparse
+
+from python.frontend.commands.args import MODEL_PATH
+from python.frontend.commands.base import BaseCommand
+
+
+class ModelCheckCommand(BaseCommand):
+    """Check if a model is supported for quantization."""
+
+    name: ClassVar[str] = "model_check"
+    aliases: ClassVar[list[str]] = ["check"]
+    help: ClassVar[str] = "Check if the model is supported for quantization."
+
+    @classmethod
+    def configure_parser(
+        cls: type[ModelCheckCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        MODEL_PATH.add_to_parser(parser)
+
+    @classmethod
+    @BaseCommand.validate_required(MODEL_PATH)
+    @BaseCommand.validate_paths(MODEL_PATH)
+    def run(cls: type[ModelCheckCommand], args: argparse.Namespace) -> None:
+        import onnx  # noqa: PLC0415
+
+        from python.core.model_processing.onnx_quantizer.exceptions import (  # noqa: PLC0415
+            InvalidParamError,
+            UnsupportedOpError,
+        )
+        from python.core.model_processing.onnx_quantizer.onnx_op_quantizer import (  # noqa: PLC0415
+            ONNXOpQuantizer,
+        )
+
+        model = onnx.load(args.model_path)
+        quantizer = ONNXOpQuantizer()
+        try:
+            quantizer.check_model(model)
+            print(f"Model {args.model_path} is supported.")  # noqa: T201
+        except UnsupportedOpError as e:
+            msg = (
+                f"Model {args.model_path} is NOT supported: "
+                f"Unsupported operations {e.unsupported_ops}"
+            )
+            raise RuntimeError(msg) from e
+        except InvalidParamError as e:
+            msg = f"Model {args.model_path} is NOT supported: {e.message}"
+            raise RuntimeError(msg) from e

python/frontend/commands/prove.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import argparse
+
+from python.core.circuits.errors import CircuitRunError
+from python.core.utils.helper_functions import CircuitExecutionConfig, RunType
+from python.frontend.commands.args import CIRCUIT_PATH, PROOF_PATH, WITNESS_PATH
+from python.frontend.commands.base import BaseCommand
+
+
+class ProveCommand(BaseCommand):
+    """Generate proof from witness."""
+
+    name: ClassVar[str] = "prove"
+    aliases: ClassVar[list[str]] = ["prov"]
+    help: ClassVar[str] = "Generate a proof from a circuit and witness."
+
+    @classmethod
+    def configure_parser(
+        cls: type[ProveCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        CIRCUIT_PATH.add_to_parser(parser)
+        WITNESS_PATH.add_to_parser(parser, "Path to an existing witness.")
+        PROOF_PATH.add_to_parser(parser)
+
+    @classmethod
+    @BaseCommand.validate_required(CIRCUIT_PATH, WITNESS_PATH, PROOF_PATH)
+    @BaseCommand.validate_paths(CIRCUIT_PATH, WITNESS_PATH)
+    @BaseCommand.validate_parent_paths(PROOF_PATH)
+    def run(cls: type[ProveCommand], args: argparse.Namespace) -> None:
+        circuit = cls._build_circuit("cli")
+
+        try:
+            circuit.base_testing(
+                CircuitExecutionConfig(
+                    run_type=RunType.PROVE_WITNESS,
+                    circuit_path=args.circuit_path,
+                    witness_file=args.witness_path,
+                    proof_file=args.proof_path,
+                    ecc=False,
+                ),
+            )
+        except CircuitRunError as e:
+            raise RuntimeError(e) from e
+
+        print(f"[prove] wrote proof → {args.proof_path}")  # noqa: T201

python/frontend/commands/verify.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import argparse
+
+from python.core.circuits.errors import CircuitRunError
+from python.core.utils.helper_functions import CircuitExecutionConfig, RunType
+from python.frontend.commands.args import (
+    CIRCUIT_PATH,
+    INPUT_PATH,
+    OUTPUT_PATH,
+    PROOF_PATH,
+    WITNESS_PATH,
+)
+from python.frontend.commands.base import BaseCommand
+
+
+class VerifyCommand(BaseCommand):
+    """Verify a proof."""
+
+    name: ClassVar[str] = "verify"
+    aliases: ClassVar[list[str]] = ["ver"]
+    help: ClassVar[str] = "Verify a proof."
+
+    @classmethod
+    def configure_parser(
+        cls: type[VerifyCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        CIRCUIT_PATH.add_to_parser(parser)
+        INPUT_PATH.add_to_parser(parser)
+        OUTPUT_PATH.add_to_parser(parser, "Path to expected outputs JSON.")
+        WITNESS_PATH.add_to_parser(parser)
+        PROOF_PATH.add_to_parser(parser)
+
+    @classmethod
+    @BaseCommand.validate_required(
+        CIRCUIT_PATH,
+        INPUT_PATH,
+        OUTPUT_PATH,
+        WITNESS_PATH,
+        PROOF_PATH,
+    )
+    @BaseCommand.validate_paths(
+        CIRCUIT_PATH,
+        INPUT_PATH,
+        OUTPUT_PATH,
+        WITNESS_PATH,
+        PROOF_PATH,
+    )
+    def run(cls: type[VerifyCommand], args: argparse.Namespace) -> None:
+        circuit = cls._build_circuit("cli")
+
+        try:
+            circuit.base_testing(
+                CircuitExecutionConfig(
+                    run_type=RunType.GEN_VERIFY,
+                    circuit_path=args.circuit_path,
+                    input_file=args.input_path,
+                    output_file=args.output_path,
+                    witness_file=args.witness_path,
+                    proof_file=args.proof_path,
+                    ecc=False,
+                ),
+            )
+        except CircuitRunError as e:
+            raise RuntimeError(e) from e
+
+        print(  # noqa: T201
+            f"[verify] verification complete for proof → {args.proof_path}",
+        )

python/frontend/commands/witness.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import argparse
+
+from python.core.circuits.errors import CircuitRunError
+from python.core.utils.helper_functions import CircuitExecutionConfig, RunType
+from python.frontend.commands.args import (
+    CIRCUIT_PATH,
+    INPUT_PATH,
+    OUTPUT_PATH,
+    WITNESS_PATH,
+)
+from python.frontend.commands.base import BaseCommand
+
+
+class WitnessCommand(BaseCommand):
+    """Generate witness from circuit and inputs."""
+
+    name: ClassVar[str] = "witness"
+    aliases: ClassVar[list[str]] = ["wit"]
+    help: ClassVar[str] = "Generate witness using a compiled circuit."
+
+    @classmethod
+    def configure_parser(
+        cls: type[WitnessCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        CIRCUIT_PATH.add_to_parser(parser)
+        INPUT_PATH.add_to_parser(parser)
+        OUTPUT_PATH.add_to_parser(parser)
+        WITNESS_PATH.add_to_parser(parser)
+
+    @classmethod
+    @BaseCommand.validate_required(
+        CIRCUIT_PATH,
+        INPUT_PATH,
+        OUTPUT_PATH,
+        WITNESS_PATH,
+    )
+    @BaseCommand.validate_paths(CIRCUIT_PATH, INPUT_PATH)
+    @BaseCommand.validate_parent_paths(OUTPUT_PATH, WITNESS_PATH)
+    def run(cls: type[WitnessCommand], args: argparse.Namespace) -> None:
+        circuit = cls._build_circuit("cli")
+
+        try:
+            circuit.base_testing(
+                CircuitExecutionConfig(
+                    run_type=RunType.GEN_WITNESS,
+                    circuit_path=args.circuit_path,
+                    input_file=args.input_path,
+                    output_file=args.output_path,
+                    witness_file=args.witness_path,
+                ),
+            )
+        except CircuitRunError as e:
+            raise RuntimeError(e) from e
+
+        print(  # noqa: T201
+            f"[witness] wrote witness → {args.witness_path} "
+            f"and outputs → {args.output_path}",
+        )