zorch.logup_gkr.jagged_prover

Jagged LogUp-GKR prover: the materialized per-layer sumcheck.

The jagged sibling of prover.GkrLayerRound. A JaggedGkrLayer materializes only sum(row_counts) of its virtual 2^(niv+nrv) positions; every non-materialized position holds the fold-neutral fraction (n=0, d=1), whose LogUp summand eq * (lam*(n0*d1 + n1*d0) + d0*d1) collapses to just its eq weight. The sumcheck therefore runs over the materialized arrays and adds the virtual mass back in closed form: the eq weights of a full hypercube sum to the product of the bound variables' eq factors (pad_adj), so the correction per round is pad_adj - eq_sum_materialized.

Round polynomials travel in coefficient form, interpolated through {0, 1, 1/2, b}: the summand carries the current variable's eq factor, whose root b = (1-z)/(1-2z) is known to both sides, so a degree-3 round needs only the materialized evaluations at {0, 1/2} plus s(1) = claim - s(0) (Gruen, https://eprint.iacr.org/2024/108). Value-form on the natural domain would need a third materialized evaluation per round.

Variables bind LSB-first (consecutive-pair fold): a jagged layer is batch-major, so the row LSB is the in-segment pair dimension and the stride-2 fold never crosses a segment boundary once odd segments are re-padded (the same in-trace gather derivation as the circuit transition). Row variables fold first while their eq factor rides as the materialized eq_row lookup; once rows are exhausted the accumulated row-eq residual becomes the scalar eq_adj and the batch variables fold densely. The bound point is challenges reversed -- LSB-first binding makes the last challenge the MSB -- so the carry convention (MSB-first point, child selector appended last) matches the dense chain's.

Per-round shapes shrink and the gather layout changes round to round, so the driver is a host-orchestrated Python loop over plain numeric bodies, not the homogeneous zorch.sumcheck scan (see docs/reference/conventions.md).

JaggedLayerProof dataclass

One jagged GKR layer's sumcheck transcript: the batching challenge and opening claim the layer entered with (the per-layer anchors a consumer diffs first when a byte-match diverges mid-pyramid), the coefficient-form round polynomials, the bound point, and the final pair openings.

A pytree (every field is an Array, like the dense sumcheck.RoundMsg) so it can be returned across a frx.jit boundary -- the per-layer jit the chained prover wraps each round in.

point is retained for wire serialization despite being replay-derivable — LayerProof.point carries the rationale and the verifier-must-never-read rule.

Source code in zorch/logup_gkr/jagged_prover.py

@partial(
    frx.tree_util.register_dataclass,
    data_fields=[
        "lam",
        "claim",
        "round_polys",
        "point",
        "numerator_0",
        "numerator_1",
        "denominator_0",
        "denominator_1",
    ],
    meta_fields=[],
)
@dataclass(frozen=True)
class JaggedLayerProof:
    """One jagged GKR layer's sumcheck transcript: the batching challenge and
    opening claim the layer entered with (the per-layer anchors a consumer
    diffs first when a byte-match diverges mid-pyramid), the coefficient-form
    round polynomials, the bound point, and the final pair openings.

    A pytree (every field is an `Array`, like the dense `sumcheck.RoundMsg`) so
    it can be returned across a `frx.jit` boundary -- the per-layer jit the
    chained prover wraps each round in.

    `point` is retained for wire serialization despite being replay-derivable
    — `LayerProof.point` carries the rationale and the
    verifier-must-never-read rule."""

    lam: Array
    claim: Array
    round_polys: Array  # (num_variables, _DEGREE + 1), ascending coefficients
    point: Array  # the bound point, MSB-first (the sampled challenges reversed)
    numerator_0: Array
    numerator_1: Array
    denominator_0: Array
    denominator_1: Array

JaggedGkrLayerRound

Bases: ProverRound

Prove one jagged GKR layer; the chain of these (floor outward) is the jagged GKR prover, threading the same (num_eval, den_eval, eval_point) carry as the dense chain. One ChallengePolicy configures lam, every per-variable fold, and the child selector; prover.bind_output consumes that same policy at the chain head.

The per-layer prove dispatches through the module-level _jagged_round_zone (one executable per layer): the schedule rides as a traced operand and the planes arrive already cap-width, so the whole-layer trace keys on the cap shape -- one compile per nrv class, reused across every layer, pass, AND shard. The round holds only its layer (no per-instance jit, no self-closure), so the chain frees each round -- and its layer -- the moment it builds the next. The pyramid stays a host-orchestrated Python loop of these (one trace per layer shape, never one jit over the whole pyramid -- it does not fit at scale; see prover.LogupSumcheckRound).

Source code in zorch/logup_gkr/jagged_prover.py

class JaggedGkrLayerRound(ProverRound):
    """Prove one jagged GKR layer; the chain of these (floor outward) is the
    jagged GKR prover, threading the same `(num_eval, den_eval, eval_point)`
    carry as the dense chain. One `ChallengePolicy` configures lam, every
    per-variable fold, and the child selector; `prover.bind_output` consumes
    that same policy at the chain head.

    The per-layer prove dispatches through the module-level `_jagged_round_zone`
    (one executable per layer): the schedule rides as a traced operand and the
    planes arrive already cap-width, so the whole-layer trace keys on the cap
    shape -- one compile per nrv class, reused across every layer, pass, AND
    shard. The round holds only its layer (no per-instance jit, no self-closure),
    so the chain frees each round -- and its layer -- the moment it builds the
    next. The pyramid stays a host-orchestrated Python loop of these (one trace
    per layer shape, never one `jit` over the whole pyramid -- it does not fit at
    scale; see `prover.LogupSumcheckRound`).
    """

    def __init__(
        self,
        layer: JaggedGkrLayer,
        challenges: ChallengePolicy,
        *,
        caps: RoundWidthCaps | None = None,
        layer_bufs: LayerBuffers | None = None,
    ) -> None:
        # `partial` closes over the args, not `self`, so the chain frees the
        # round -- and its layer -- the moment it builds the next. Pass ONE
        # `layer_bufs` per chain (None materializes the cap pad fresh).
        policy = challenges
        challenge_limbs = policy.base_limbs
        self._call = partial(
            _jagged_round_via_zone, layer, challenge_limbs, caps, layer_bufs
        )

    def __call__(
        self, claim: LayerClaim, transcript: Transcript
    ) -> tuple[LayerClaim, Transcript, JaggedLayerProof]:
        return self._call(claim, transcript)

prove_jagged_layer

prove_jagged_layer(
    layer: JaggedGkrLayer,
    lam: Array,
    claim: Array,
    eval_point: Array,
    transcript: Transcript,
    *,
    challenges: ChallengePolicy,
    caps: RoundWidthCaps | None = None
) -> tuple[Array, Transcript, JaggedLayerProof]

Prove one jagged GKR layer's materialized sumcheck from an explicit lam / claim (no inter-layer carry). The standalone single-layer seam the layer tests drive; the pyramid runs JaggedGkrLayerRound, which brackets this same core (_prove_jagged_layer_from_ops) with the carry.

eval_point is MSB-first over (batch || row) variables; its length fixes the virtual row depth nrv = len(eval_point) - niv, which may exceed what the materialized row counts need -- the extra rounds fold saturated all-ones segments against re-padded neutral rows, exactly the virtual positions' values. Returns the bound point (MSB-first, i.e. the challenges reversed), the advanced transcript, and the proof.

caps is the fixed-width, size-invariant round layout, and is mandatory: every round runs at one static operand shape per phase, live prefix tracked by the rounds' live operand, so one compiled round kernel serves every round -- and every layer and input proved under the same caps. Row counts are traced, so there is no exact-layout fallback with static per-round output widths; the zero-slack layout is the capacity layout whose caps happen to be tight.

The device-derived schedule: the per-round re-pad schedule is a pure function of row_counts + the round index and derives inside the claimed kernels, so the loop carries only the tiny i32[nseg] row_counts operand plus per-round i32[3] live triples -- both ride as traced operands, so row_counts never enters the jit key. The virtual-row-space fit (max(row_counts) <= 2^nrv) is the consumer's capacity-class obligation -- a host check cannot read the traced counts.

Source code in zorch/logup_gkr/jagged_prover.py

def prove_jagged_layer(
    layer: JaggedGkrLayer,
    lam: Array,
    claim: Array,
    eval_point: Array,
    transcript: Transcript,
    *,
    challenges: ChallengePolicy,
    caps: RoundWidthCaps | None = None,
) -> tuple[Array, Transcript, JaggedLayerProof]:
    """Prove one jagged GKR layer's materialized sumcheck from an explicit
    `lam` / `claim` (no inter-layer carry). The standalone single-layer seam
    the layer tests drive; the pyramid runs `JaggedGkrLayerRound`, which
    brackets this same core (`_prove_jagged_layer_from_ops`) with the carry.

    `eval_point` is MSB-first over (batch || row) variables; its length
    fixes the virtual row depth `nrv = len(eval_point) - niv`, which may
    exceed what the materialized row counts need -- the extra rounds fold
    saturated all-ones segments against re-padded neutral rows, exactly the
    virtual positions' values. Returns the bound point (MSB-first, i.e. the
    challenges reversed), the advanced transcript, and the proof.

    `caps` is the fixed-width, size-invariant round layout, and is
    mandatory: every round runs at one static operand shape per phase, live
    prefix tracked by the rounds' `live` operand, so one compiled round
    kernel serves every round -- and every layer and input proved under the
    same caps. Row counts are traced, so there is no exact-layout fallback
    with static per-round output widths; the zero-slack layout is the capacity
    layout whose caps happen to be tight.

    The device-derived schedule: the per-round re-pad schedule is a
    pure function of `row_counts` + the round index and derives inside the
    claimed kernels, so the loop carries only the tiny i32[nseg] `row_counts`
    operand plus per-round i32[3] live triples -- both ride as traced operands,
    so `row_counts` never enters the jit key. The virtual-row-space fit
    (`max(row_counts) <= 2^nrv`) is the consumer's capacity-class obligation
    -- a host check cannot read the traced counts.
    """
    if caps is None:
        raise ValueError("a jagged layer proves under caps; pass RoundWidthCaps")
    challenge_limbs = challenges.limbs_over(transcript.field)
    niv = layer.num_batch_variables
    nrv = eval_point.shape[0] - niv
    if nrv < 1:
        raise ValueError(
            f"eval_point must carry at least one row variable: got "
            f"{eval_point.shape[0]} coordinates for {niv} batch variables"
        )
    planes = _Planes(
        layer.numerator_0,
        layer.numerator_1,
        layer.denominator_0,
        layer.denominator_1,
    )
    return _prove_jagged_layer_from_ops(
        planes,
        niv,
        layer.row_counts,
        _derive_live_meta(layer.row_counts, nrv),
        None,
        lam,
        claim,
        eval_point,
        transcript,
        challenge_limbs,
        caps,
    )