# Sumcheck — the workhorse interactive proof behind modern SNARKs.
#
# The prover convinces a verifier of the value of  sum over x in {0,1}^n  of
# a(x)·b(x)  — the whole 2^n-entry hypercube — while the verifier does only
# O(n) work. Fiat-Shamir runs over a real poseidon2 (koalabear-16) transcript.
# Everything here is plain zorch; change LOG_N (10..24) or the seeds and re-run.
import jax.numpy as jnp
from zk_dtypes import koalabear_mont as F

from zorch.hash.poseidon2.testing.koalabear16 import koalabear16_perm
from zorch.prove import fold_rounds
from zorch.sumcheck import prover, verifier
from zorch.testkit.random_field import rand_field
from zorch.transcript import DuplexTranscript
from zorch.verify import verify

LOG_N = 20  # 2^20 = 1,048,576 evaluations


def transcript():
    return DuplexTranscript.new(koalabear16_perm(), rate=8)


a = rand_field(1, (1 << LOG_N,), F)
b = rand_field(2, (1 << LOG_N,), F)
claimed = jnp.sum(a * b)

state = jnp.stack([a, b])
_, _, msgs = fold_rounds(
    prover.StandardRound(prover.ProductSummand(2)), state, transcript(), LOG_N
)
proof = jnp.stack(msgs)

point, final_claim, _, ok = verify(
    verifier.SumcheckRound(2), claimed, proof, transcript()
)
print("verifier accepts:", bool(ok))
print(f"proof: {proof.shape[0]} rounds x degree-{proof.shape[1] - 1} polynomials")
print(f"claimed sum: {int(claimed)}")

zorch

SNARK = Σ IOP Round

FRX-native building blocks for Modern SNARKs. zorch sits between FRX — Fractalyze's fork of JAX — and the proof systems that consume it: FRX provides tracing and codegen, lowered through Fractalyze XLA, its fork of stock XLA that adds native field and elliptic-curve types. zorch provides the reusable pieces a proof system is assembled from.

A Modern SNARK is IOP + PCS. The way deep learning stacks Layers, zorch stacks Rounds — the one composable unit the rest is threaded through.

Design Philosophy

• Proving-scheme-agnostic. The blocks capture every proving scheme, not a single one. Round / Fiat-Shamir / Polynomial / PCS / fold / zero-check compose into FRI, sumcheck, GKR, STARK, Basefold, WHIR, …; pairing-based schemes plug in by swapping the PCS block (e.g. a KZG-style commitment).
• Implementation-agnostic. zorch targets the proving scheme, not any one downstream implementation — a zkVM, a zkML prover, a zkTLS prover. Each plugs in as a consumer; nothing implementation-specific leaks into a block.
• Fusion-first. Each Round — and each commit/open, each absorb/squeeze, each fold step, and a hash permutation's internal rounds — must lower to a single fused kernel. We get there by construction, not by a per-primitive pattern-matcher in the compiler.
• Easy to assemble. These are building blocks; the API optimizes for snapping them together.

Building blocks

The one unit is the Round — a prover↔verifier interaction of an IOP: a message observed into the Fiat-Shamir transcript, a challenge sampled back (observe→sample, via __call__). Rounds nest — a single per-variable step is a Round, and a whole sumcheck (its per-variable Rounds bundled) is itself a Round. This is what SNARK = Σ IOP Round says literally: a Fiat-Shamir-compiled IOP is a tree of these rounds, Σ flattening it to the leaf interactions.

Grouping Rounds gives either a bigger Round (a sumcheck, from its per-variable rounds) or — when the group is a top-level phase — a Stage. Two roles organize the composition; both are Rounds, since a chain is itself one:

• A Stage is a Round that is one phase of the scheme's prove_chain — the sequence of Stages the scheme is (trace-commit, logup-gkr, zero-check, a PCS opening).
• A Bridge is a transcript-only Round inside a Stage — a grind, a framed observe, a sampled-and-discarded challenge — soundness or security work the phase needs, not a phase of its own: a grind buys security bits, framing / domain separation closes a Fiat-Shamir soundness gap, a discarded sample matches the reference's schedule.

So the shape is recursive — the prove_chain is a sequence of Stages; a Stage chains Rounds and Bridges; a Round may itself chain Rounds, down to the leaf interaction:

prove()  —  the prove_chain is Stages; a Stage holds Rounds and Bridges
──────────────────────────────────────────────────────────────────────

  Stage   trace-commit          commit the witness columns
  Stage   logup-gkr             the interaction argument:
    Bridge  grind                 a PoW inside the stage (buys security bits)
    Round   layer L                one layer — itself a Round of Rounds:
      Round   bind x₀                a leaf: one observe → sample
      Round   bind x₁
    Round   layer L-1
  Stage   zero-check            the constraint sumcheck:
    Bridge  observe(framing)      bind the transcript first (soundness)
    Round   bind x₀                a leaf: one observe → sample
    Round   bind x₁
  Stage   jagged-evals          the PCS opening

Stage and Bridge are the same Round interface — that is how chains nest and how the verifier mirrors the prover round-for-round — but the roles are what a reader navigates by.

Where the boundaries fall. A leaf Round is each prover↔verifier interaction (observe→sample); Rounds bundle into a bigger Round or, at a prove_chain phase, a Stage; a Bridge sits inside a Stage wherever the reference's soundness argument needs a transcript op. The full carry-and-seam contract is docs/composition/stage-composition.md.

Where the classic pieces fit. A ZK reader expects Fiat-Shamir, Polynomial, PCS, and sumcheck as top-level "blocks." In this picture they are not peers of the Round:

• Fiat-Shamir, Polynomial, and fold are the materials a Round body computes with — the transcript it threads, the polynomials it evaluates, the fold (2-to-1 reduction, one challenge per round) it applies each step.
• A PCS opening and a zero-check are Stages — each a distinct phase: a zero-check reduces to a sumcheck, while a PCS opening runs its commitment-opening and evaluation checks (the jagged-evals stage above).

Development

zorch is pure Python on FRX, run against its GPU plugin. A virtualenv with the pinned toolchain:

python3.11 -m venv .venv && . .venv/bin/activate
pip install -r requirements.in \
    --extra-index-url https://fractalyze.github.io/pypi/simple/

The dev loop — per-workspace venvs, developing against a local Fractalyze XLA build, the FRX compile-cache rule — lives in docs/reference/development.md.

Documentation

• Task-indexed docs hub: docs/README.md — indexes every design doc by what you're trying to do.

License

Licensed under the Apache License, Version 2.0 (see LICENSE).