JSTprove 1.0.0__py3-none-macosx_11_0_arm64.whl

python/core/utils/model_registry.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import importlib
+import pkgutil
+from collections import namedtuple
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import types
+    from collections.abc import Callable
+
+
+from python.core import circuit_models, circuits
+from python.core.circuit_models.generic_onnx import GenericModelONNX
+from python.core.circuits.base import Circuit
+
+ModelEntry = namedtuple(  # noqa: PYI024
+    "ModelEntry",
+    ["name", "source", "loader", "args", "kwargs"],
+)
+
+
+def scan_model_files(
+    directory: str,
+    extension: str,
+    loader_fn: Callable,
+    prefix: str,
+) -> list[ModelEntry]:
+    """Scan a directory for model files and return each discovered model in a callable.
+
+    Args:
+        directory (str): Path to the directory to scan.
+        extension (str): File extension to filter by (e.g., ".onnx").
+        loader_fn (Callable):
+            Loader function that can instantiate a model from a file path.
+        prefix (str): Source prefix used to categorize models (e.g., "onnx").
+
+    Returns:
+        list[ModelEntry]: A list of ModelEntry objects representing discovered models.
+    """
+    _ = extension
+    entries = []
+    for item in Path(directory).iterdir():
+        file_or_foldername = item.name
+        if prefix == "onnx" and (
+            item.is_file() and file_or_foldername.endswith(".onnx")
+        ):
+            name = file_or_foldername[:-5]
+            path = str(item)
+            entries.append(
+                ModelEntry(
+                    name=f"{name}",
+                    source=prefix,
+                    loader=lambda p=path: loader_fn(p),
+                    args=(),
+                    kwargs={},
+                ),
+            )
+    return entries
+
+
+def import_all_submodules(package: types.ModuleType) -> None:
+    """Import all submodules of a given package.
+
+    Args:
+        package (types.ModuleType): The Python package to import submodules from.
+    """
+    for _, name, _ in pkgutil.walk_packages(
+        package.__path__,
+        package.__name__ + ".",
+    ):
+        importlib.import_module(name)
+
+
+# Import all submodules so their classes are registered
+import_all_submodules(circuit_models)
+import_all_submodules(circuits)
+
+
+def onnx_test_loader(path: str) -> GenericModelONNX:
+    return GenericModelONNX(path, use_find_model=True)
+
+
+def all_subclasses(cls) -> set:  # noqa: ANN001
+    """Recursively find all subclasses of a given class.
+
+    Args:
+        cls (type): The base class to search subclasses for.
+
+    Returns:
+        _type_: A set of all subclasses of the given class.
+    """
+    subclasses = set(cls.__subclasses__())
+    return subclasses.union(s for c in subclasses for s in all_subclasses(c))
+
+
+def build_models_to_test() -> list[ModelEntry]:
+    """Build a list of model entries to be tested.
+
+    - Collects all subclasses of the Circuit base class.
+    - Filters out unwanted base or placeholder models.
+    - Adds discovered ONNX models from the models directory.
+
+    Returns:
+        list[ModelEntry]: A list of model entries to be used in tests.
+    """
+    models = []
+    for cls in all_subclasses(Circuit):
+        name = cls.__name__.lower()
+        models.append(
+            ModelEntry(name=name, source="class", loader=cls, args=(), kwargs={}),
+        )
+    # Filter unwanted class models
+    models = [
+        m
+        for m in models
+        if m.name not in {"zkmodel", "genericmodelonnx", "zktorchmodel", "zkmodelbase"}
+    ]
+    # Add ONNX models
+    models += scan_model_files(
+        "python/models/models_onnx",
+        ".onnx",
+        onnx_test_loader,
+        "onnx",
+    )
+    return models
+
+
+MODELS_TO_TEST = build_models_to_test()
+
+
+def list_available_models() -> list[str]:
+    """list all available models in a human-readable format.
+
+    Returns:
+        list[str]: A sorted list of strings in the form "source: model_name".
+    """
+    return sorted(f"{model.source}: {model.name}" for model in MODELS_TO_TEST)
+
+
+def get_models_to_test(
+    selected_models: list[str] | None = None,
+    source_filter: str | None = None,
+) -> list[ModelEntry]:
+    """Retrieve models to be tested with optional filtering.
+
+    Args:
+        selected_models (list[str], optional):
+            A list of model names to include. Defaults to None.
+        source_filter (str, optional):
+            Restrict models to a specific source (e.g., "onnx", "class").
+            Defaults to None.
+
+    Returns:
+        list[ModelEntry]: A filtered list of model entries.
+    """
+    models = MODELS_TO_TEST
+
+    if selected_models is not None:
+        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]
+
+    return models

python/core/utils/scratch_tests.py

@@ -0,0 +1,66 @@
+import onnx
+from onnx import TensorProto, helper, shape_inference
+from onnx import numpy_helper
+from onnx import load, save
+from onnx.utils import extract_model
+
+def prune_model(model_path, output_names, save_path):
+    model = load(model_path)
+
+    # 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
+    )
+
+    print(f"Pruned model saved to {save_path}")
+
+
+def cut_model(model_path, output_names, save_path):
+    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:
+        graph.output.pop()
+
+    # Add new outputs
+    for name in output_names:
+        # 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.")
+
+        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)
+        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}")
+
+
+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"
+    )

python/core/utils/witness_utils.py

@@ -0,0 +1,291 @@
+from __future__ import annotations
+
+import struct
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING, BinaryIO
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+from python.core.utils.errors import ProofSystemNotImplementedError
+from python.core.utils.helper_functions import ZKProofSystems
+
+
+# -------------------------
+# Base Witness Loader
+# -------------------------
+class WitnessLoader(ABC):
+    def __init__(self: WitnessLoader, path: str) -> None:
+        self.path = path
+
+    @abstractmethod
+    def load_witness(self: WitnessLoader) -> dict:
+        """Load witness data from file."""
+
+    @abstractmethod
+    def compare_witness_to_io(
+        self: WitnessLoader,
+        witnesses: dict,
+        expected_inputs: dict,
+        expected_outputs: dict,
+        modulus: int,
+    ) -> bool:
+        """Compare witness to expected I/O."""
+
+
+def read_usize(f: BinaryIO, usize_len: int | None = 8) -> int:
+    """
+    Read an unsigned integer of size `usize_len` bytes from a binary file object.
+
+    Args:
+        f (BinaryIO): Opened file in binary mode.
+        usize_len (int, optional): Number of bytes to read
+        (default is 8, for 64-bit systems).
+
+    Returns:
+        int: The unpacked unsigned integer.
+    """
+    return struct.unpack("<Q", f.read(usize_len))[0]
+
+
+def read_u256(f: BinaryIO) -> int:
+    """
+    Read a 256-bit unsigned integer (U256) from a binary file in little-endian format.
+
+    Args:
+        f (BinaryIO): Opened file in binary mode.
+
+    Returns:
+        int: The 256-bit integer.
+    """
+    return int.from_bytes(f.read(32), "little")
+
+
+def read_field_elements(f: BinaryIO, count: int) -> list[int]:
+    """
+    Read a sequence of 32-byte field elements from a binary file.
+
+    Args:
+        f (BinaryIO): Opened file in binary mode.
+        count (int): Number of 32-byte elements to read.
+
+    Returns:
+        list[int]: List of integers representing the field elements.
+    """
+    return [read_u256(f) for _ in range(count)]
+
+
+def to_field_repr(value: int, modulus: int) -> int:
+    """
+    Convert a signed integer to its field representation modulo `modulus`.
+
+    Args:
+        value (int): Integer to convert.
+        modulus (int): Field modulus.
+
+    Returns:
+        int: Least field representation of the integer, ensuring a non-negative result.
+    """
+    return value % modulus
+
+
+class ExpanderWitnessLoader(WitnessLoader):
+    def load_witness(self: ExpanderWitnessLoader) -> dict:
+        """
+        Load witness data from a binary file and return it in structured form.
+
+        Returns:
+            dict: Dictionary containing:
+                - num_witnesses (int)
+                - num_inputs_per_witness (int)
+                - num_public_inputs_per_witness (int)
+                - modulus (int)
+                - witnesses (list of dicts with 'inputs' and 'public_inputs')
+        """
+        path = self.path
+        with Path(path).open("rb") as f:
+            num_witnesses = read_usize(f)
+            num_inputs = read_usize(f)
+            num_public_inputs = read_usize(f)
+            modulus = read_u256(f)
+
+            total = num_witnesses * (num_inputs + num_public_inputs)
+            values = read_field_elements(f, total)
+
+        # Reshape into witnesses
+        witnesses = []
+        offset = 0
+        for _ in range(num_witnesses):
+            inputs = values[offset : offset + num_inputs]
+            public_inputs = values[
+                offset + num_inputs : offset + num_inputs + num_public_inputs
+            ]
+            witnesses.append({"inputs": inputs, "public_inputs": public_inputs})
+            offset += num_inputs + num_public_inputs
+
+        return {
+            "num_witnesses": num_witnesses,
+            "num_inputs_per_witness": num_inputs,
+            "num_public_inputs_per_witness": num_public_inputs,
+            "modulus": modulus,
+            "witnesses": witnesses,
+        }
+
+    def compare_witness_to_io(
+        self: ExpanderWitnessLoader,
+        witnesses: dict,
+        expected_inputs: dict,
+        expected_outputs: dict,
+        modulus: int,
+        scaling_function: Callable[[list[int], int, int], list[int]] | None = None,
+    ) -> bool:
+        """
+        Compare the public inputs of the first witness
+        against expected inputs and outputs.
+
+        Accounts for negative numbers by representing them in the field as
+        `modulus - abs(value)`.
+
+        Args:
+            witnesses (dict): Witness data as returned by `load_witness`.
+            expected_inputs (dict):
+                Dictionary containing key "input"
+                mapping to a list of expected input integers.
+            expected_outputs (dict):
+                Dictionary containing key "output"
+                mapping to a list of expected output integers.
+            modulus (int): Field modulus.
+            scaling_function
+                (Callable[[list[int], int, int], list[int]] | None, optional):
+                    Optional scaling function to apply to inputs.
+                    Takes (inputs, scale_base, scale_exponent)
+                    and returns scaled inputs. Defaults to None.
+
+        Returns:
+            bool:
+            True if the witness public inputs match the expected inputs and outputs,
+            False otherwise.
+        """
+
+        import torch  # noqa: PLC0415
+
+        scale_base = witnesses["witnesses"][0]["public_inputs"][-2]
+        scale_exponent = witnesses["witnesses"][0]["public_inputs"][-1]
+
+        # Convert expectations into field form
+        inputs_list = expected_inputs.get("input", [])
+
+        if callable(scaling_function):
+            inputs_list = (
+                torch.tensor(scaling_function(inputs_list, scale_base, scale_exponent))
+                .flatten()
+                .tolist()
+            )
+        else:
+            inputs_list = (
+                torch.round(
+                    torch.tensor(inputs_list) * (scale_base**scale_exponent),
+                )
+                .long()
+                .tolist()
+            )
+        outputs_list = expected_outputs.get("output", [])
+
+        expected_inputs_mod = [
+            to_field_repr(v, modulus)
+            for v in torch.tensor(inputs_list).flatten().tolist()
+        ]
+        expected_outputs_mod = [to_field_repr(v, modulus) for v in outputs_list]
+
+        n_inputs = len(expected_inputs_mod) + len(expected_outputs_mod) + 2
+
+        witness = witnesses["witnesses"][0]["public_inputs"]
+
+        if n_inputs != len(witness):
+            return False
+
+        actual_inputs = witness[: len(expected_inputs_mod)]
+        actual_outputs = witness[len(expected_inputs_mod) : -2]
+
+        # Compare
+        return (
+            actual_inputs == expected_inputs_mod
+            and actual_outputs == expected_outputs_mod
+        )
+
+
+# -------------------------
+# Factory
+# -------------------------
+def get_loader(system: ZKProofSystems, path: str) -> WitnessLoader:
+    if system == ZKProofSystems.Expander:
+        return ExpanderWitnessLoader(path)
+    msg = f"No loader implemented for {system}"
+    raise ProofSystemNotImplementedError(msg)
+
+
+# -------------------------
+# Public API
+# -------------------------
+def load_witness(path: str, system: ZKProofSystems = ZKProofSystems.Expander) -> dict:
+    loader = get_loader(system, path)
+    return loader.load_witness()
+
+
+def compare_witness_to_io(  # noqa: PLR0913
+    witnesses: dict,
+    expected_inputs: dict,
+    expected_outputs: dict,
+    modulus: int,
+    system: ZKProofSystems = ZKProofSystems.Expander,
+    scaling_function: Callable[[list[int], int, int], list[int]] | None = None,
+) -> bool:
+    """
+    Compare witness data to expected inputs and outputs for a given ZK proof system.
+
+    Args:
+        witnesses (dict): Witness data as returned by `load_witness`.
+        expected_inputs (dict):
+            Dictionary containing key "input" mapping to list of expected integers.
+        expected_outputs (dict):
+            Dictionary containing key "output" mapping to list of expected integers.
+        modulus (int): Field modulus.
+        system (ZKProofSystems, optional):
+            The ZK proof system. Defaults to ZKProofSystems.Expander.
+        scaling_function (Callable[[list[int], int, int], list[int]] | None, optional):
+            Optional scaling function to apply to inputs.
+            Takes (inputs, scale_base, scale_exponent) and returns scaled inputs.
+            Defaults to None.
+
+    Returns:
+        bool: True if the witness matches the expected I/O, False otherwise.
+    """
+    loader = get_loader(system, "")  # path not needed for comparison
+    return loader.compare_witness_to_io(
+        witnesses,
+        expected_inputs,
+        expected_outputs,
+        modulus,
+        scaling_function,
+    )
+
+
+if __name__ == "__main__":
+    import time
+
+    start_time = time.time()
+    w = load_witness("./artifacts/lenet/witness.bin", ZKProofSystems.Expander)
+    end_time = time.time()
+    print("Modulus:", w["modulus"])  # noqa: T201
+    print("First witness inputs:", w["witnesses"][0]["inputs"][0])  # noqa: T201
+    print(  # noqa: T201
+        "First witness public inputs:",
+        w["witnesses"][0]["public_inputs"][0],
+    )
+
+    print(len(w["witnesses"][0]["public_inputs"]))  # noqa: T201
+    print((w["witnesses"][0]["public_inputs"][0] - w["modulus"]) / 2**18)  # noqa: T201
+    elapsed = end_time - start_time
+
+    print("time taken: ", elapsed)  # noqa: T201

python/frontend/cli.py

@@ -0,0 +1,115 @@
+# python/frontend/cli.py
+"""JSTprove CLI."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from python.frontend.commands import BaseCommand
+
+from python.frontend.commands import (
+    BenchCommand,
+    CompileCommand,
+    ModelCheckCommand,
+    ProveCommand,
+    VerifyCommand,
+    WitnessCommand,
+)
+from python.frontend.commands.base import HiddenPositionalHelpFormatter
+
+BANNER_TITLE = r"""
+  888888  .d8888b. 88888888888
+    "88b d88P  Y88b    888
+     888 Y88b.         888
+     888  "Y888b.      888  88888b.  888d888 .d88b.  888  888  .d88b.
+     888     "Y88b.    888  888 "88b 888P"  d88""88b 888  888 d8P  Y8b
+     888       "888    888  888  888 888    888  888 Y88  88P 88888888
+     88P Y88b  d88P    888  888 d88P 888    Y88..88P  Y8bd8P  Y8b.
+     888  "Y8888P"     888  88888P"  888     "Y88P"    Y88P    "Y8888
+   .d88P                    888
+ .d88P"                     888
+888P"                       888
+"""
+
+COMMANDS: list[type[BaseCommand]] = [
+    ModelCheckCommand,
+    CompileCommand,
+    WitnessCommand,
+    ProveCommand,
+    VerifyCommand,
+    BenchCommand,
+]
+
+
+def print_header() -> None:
+    """Print the CLI banner (no side-effects at import time)."""
+    print(  # noqa: T201
+        BANNER_TITLE
+        + "\n"
+        + "JSTprove — Verifiable ML by Inference Labs\n"
+        + "Based on Polyhedra Network's Expander (GKR-based proving system)\n",
+    )
+
+
+def main(argv: list[str] | None = None) -> int:
+    """
+    Entry point for the JSTprove CLI.
+
+    Returns:
+      0 on success, 1 on error.
+    """
+    argv = sys.argv[1:] if argv is None else argv
+
+    parser = argparse.ArgumentParser(
+        prog="jst",
+        description="ZKML CLI (compile, witness, prove, verify).",
+        allow_abbrev=False,
+    )
+    parser.add_argument(
+        "--no-banner",
+        action="store_true",
+        help="Suppress the startup banner.",
+    )
+
+    subparsers = parser.add_subparsers(dest="cmd", required=True)
+
+    command_map = {}
+    for command_cls in COMMANDS:
+        cmd_parser = subparsers.add_parser(
+            command_cls.name,
+            aliases=command_cls.aliases,
+            help=command_cls.help,
+            allow_abbrev=False,
+            formatter_class=HiddenPositionalHelpFormatter,
+        )
+        command_cls.configure_parser(cmd_parser)
+        command_map[command_cls.name] = command_cls
+        for alias in command_cls.aliases:
+            command_map[alias] = command_cls
+
+    args = parser.parse_args(argv)
+
+    if not args.no_banner and not os.environ.get("JSTPROVE_NO_BANNER"):
+        print_header()
+
+    try:
+        command_cls = command_map[args.cmd]
+        command_cls.run(args)
+    except (ValueError, FileNotFoundError, PermissionError, RuntimeError) as e:
+        print(f"Error: {e}", file=sys.stderr)  # noqa: T201
+        return 1
+    except SystemExit:
+        raise
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)  # noqa: T201
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

python/frontend/commands/__init__.py

@@ -0,0 +1,17 @@
+from python.frontend.commands.base import BaseCommand
+from python.frontend.commands.bench import BenchCommand
+from python.frontend.commands.compile import CompileCommand
+from python.frontend.commands.model_check import ModelCheckCommand
+from python.frontend.commands.prove import ProveCommand
+from python.frontend.commands.verify import VerifyCommand
+from python.frontend.commands.witness import WitnessCommand
+
+__all__ = [
+    "BaseCommand",
+    "BenchCommand",
+    "CompileCommand",
+    "ModelCheckCommand",
+    "ProveCommand",
+    "VerifyCommand",
+    "WitnessCommand",
+]