Skip to content

zorch.logup_gkr.jagged_stage

Jagged LogUp-GKR role implementations.

The same reduction the dense roles run -- a circuit-output claim down to an input-layer point claim -- over the jagged layout. Claim types come from stage, so the two layouts are interchangeable at the stage seam and only the witness and the layer proofs differ.

JaggedGkrWitness dataclass

The input jagged layer and the per-transition fold schedule carrying it to the batch floor.

The schedule follows the row counts, so it is per-input and rides the witness; the width caps are a capacity class shared across inputs and ride the prover instead. schedules[k] is the argument circuit.jagged_layer_transition takes -- (out_row_counts, out_width) traced, or a bare host sequence at its zero-slack width.

Source code in zorch/logup_gkr/jagged_stage.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
@dataclass(frozen=True)
class JaggedGkrWitness:
    """The input jagged layer and the per-transition fold schedule carrying it
    to the batch floor.

    The schedule follows the row counts, so it is per-input and rides the
    witness; the width caps are a capacity class shared across inputs and ride
    the prover instead. `schedules[k]` is the argument
    `circuit.jagged_layer_transition` takes -- `(out_row_counts, out_width)`
    traced, or a bare host sequence at its zero-slack width.
    """

    input_layer: JaggedGkrLayer
    schedules: Sequence[tuple[Array, int] | Sequence[int]]

JaggedLogUpGkrProver

Bases: ProverStage[LogUpOutputClaim, JaggedGkrWitness, InputLayerClaim, GkrProof[JaggedLayerProof]]

Prove an output claim conditional on an input-layer claim, jagged.

Source code in zorch/logup_gkr/jagged_stage.py 
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
class JaggedLogUpGkrProver(
    ProverStage[
        LogUpOutputClaim,
        JaggedGkrWitness,
        InputLayerClaim,
        GkrProof[JaggedLayerProof],
    ]
):
    """Prove an output claim conditional on an input-layer claim, jagged."""

    def __init__(
        self,
        caps: RoundWidthCaps,
        challenges: ChallengePolicy,
    ) -> None:
        self.caps = caps
        self.challenges = challenges

    def prove(
        self,
        claim: LogUpOutputClaim,
        witness: JaggedGkrWitness,
        transcript: Transcript,
    ) -> ProveResult[InputLayerClaim, GkrProof[JaggedLayerProof]]:
        if len(witness.schedules) != claim.layers:
            raise ValueError(
                f"claim expects {claim.layers} GKR layers, "
                f"witness folds through {len(witness.schedules)}"
            )
        pyramid = build_jagged_pyramid(witness.input_layer, witness.schedules)
        # The floor holds the public output; the chain proves the rest.
        pyramid.pop()
        carry, transcript = bind_output(claim.output, transcript, self.challenges)
        # One `LayerBuffers` per chain: the cap-wide planes are ~GiB per class,
        # so the holder must die with the prove. Layers leave the list as the
        # chain consumes them, keeping one intermediate layer resident instead
        # of the whole pyramid.
        buffers = LayerBuffers()

        def rounds() -> Iterator[ProverLayerRound]:
            while pyramid:
                yield ProverLayerRound(
                    pyramid.pop(),
                    self.challenges,
                    caps=self.caps,
                    layer_bufs=buffers,
                )

        carry, transcript, proofs = prove_rounds(rounds(), carry, transcript)
        return ProveResult(
            _input_claim(carry),
            GkrProof(tuple(proofs)),
            transcript,
        )

JaggedLogUpGkrVerifier

Bases: VerifierStage[LogUpOutputClaim, InputLayerClaim, GkrProof[JaggedLayerProof]]

Verify an output-to-input-layer LogUp-GKR reduction, jagged.

Layout-blind, exactly like its rounds: the prover's virtual-mass corrections make the round polynomials those of the virtual dense hypercube, so nothing here reads a row count.

Source code in zorch/logup_gkr/jagged_stage.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
class JaggedLogUpGkrVerifier(
    VerifierStage[LogUpOutputClaim, InputLayerClaim, GkrProof[JaggedLayerProof]]
):
    """Verify an output-to-input-layer LogUp-GKR reduction, jagged.

    Layout-blind, exactly like its rounds: the prover's virtual-mass
    corrections make the round polynomials those of the virtual dense
    hypercube, so nothing here reads a row count.
    """

    def __init__(self, challenges: ChallengePolicy) -> None:
        self.challenges = challenges

    def verify(
        self,
        claim: LogUpOutputClaim,
        reduction_proof: GkrProof[JaggedLayerProof],
        transcript: Transcript,
    ) -> VerifyResult[InputLayerClaim]:
        if len(reduction_proof.layers) != claim.layers:
            raise ValueError(
                f"expected {claim.layers} GKR layers, "
                f"got {len(reduction_proof.layers)}"
            )
        carry, transcript = bind_output(claim.output, transcript, self.challenges)
        carry, transcript, ok = verify_rounds(
            (VerifierLayerRound(self.challenges) for _ in range(claim.layers)),
            carry,
            reduction_proof.layers,
            transcript,
        )
        return VerifyResult(_input_claim(carry), transcript, ok)