zorch.logup_gkr.circuit

LogUp-GKR fractional-sum circuit.

A layer is four equal-length MLEs (n0, n1, d0, d1): index i carries two binary-tree children, the fractions n0[i]/d0[i] and n1[i]/d1[i]. A transition folds each pair into one fraction --

n0/d0 + n1/d1 = (n0*d1 + n1*d0) / (d0*d1)

-- pairing LSB-adjacent nodes (stride 2), which halves every MLE and eliminates one (row) variable. Iterating to the batch floor (num_row_variables == 0) leaves one fraction per batch element; extract_outputs interleaves the two children back into the output numerator/denominator MLEs over the batch variables plus one.

The four MLEs share a length, but the numerator pair and the denominator pair need not share a field: a base-field numerator under an extension-field denominator promotes to their common field at the fold (n0*d1 + n1*d0), so a consumer can keep a layer's numerator reads narrow until its first transition. zorch stays scheme-agnostic about why a consumer would; it only guarantees the promotion is byte-identical to folding an all-extension copy.

Two layouts share that fold. Dense (GkrLayer): every batch element has one row count, so a layer is a flat power of two with no padding. Jagged (JaggedGkrLayer): the MLEs are stored batch-major with per-batch-element row counts; a transition pre-pads odd segments and post-pads to a consumer-supplied schedule with the additive-identity fraction (n=0, d=1), so the flat stride-2 fold never pairs across a batch boundary. Which heights the schedule pads to is the consumer's policy, and interaction fingerprinting likewise stays in the consumer -- zorch stays scheme-agnostic.

GkrLayer dataclass

One dense fractional-sum layer over (batch || row) variables.

num_batch_variables is the floor: folding stops once the row variables are exhausted and only the batch dimension remains. Each batch element is one independent LogUp instance -- a consumer may call it a lookup interaction (the term used throughout this module's circuit prose); zorch itself stays scheme-agnostic.

Source code in zorch/logup_gkr/circuit.py

@dataclass(frozen=True)
class GkrLayer:
    """One dense fractional-sum layer over (batch || row) variables.

    `num_batch_variables` is the floor: folding stops once the row
    variables are exhausted and only the batch dimension remains. Each
    batch element is one independent LogUp instance -- a consumer may call
    it a lookup *interaction* (the term used throughout this module's
    circuit prose); zorch itself stays scheme-agnostic.
    """

    numerator_0: Array
    numerator_1: Array
    denominator_0: Array
    denominator_1: Array
    num_batch_variables: int

    def __post_init__(self) -> None:
        # Reject malformed layers at construction rather than at a later
        # broadcast or negative-row-count. num_variables also checks the
        # power-of-two width via log2_strict_usize.
        shape = self.numerator_0.shape
        for name in ("numerator_1", "denominator_0", "denominator_1"):
            if getattr(self, name).shape != shape:
                raise ValueError(
                    f"all MLEs must share a shape; {name} is "
                    f"{getattr(self, name).shape}, numerator_0 is {shape}"
                )
        if not 0 <= self.num_batch_variables <= self.num_variables:
            raise ValueError(
                f"num_batch_variables must be in [0, {self.num_variables}], "
                f"got {self.num_batch_variables}"
            )

    @property
    def num_variables(self) -> int:
        return log2_strict_usize(self.numerator_0.shape[0])

    @property
    def num_row_variables(self) -> int:
        return self.num_variables - self.num_batch_variables

LogUpGkrOutput dataclass

Final numerator/denominator MLEs after all transitions.

Source code in zorch/logup_gkr/circuit.py

@dataclass(frozen=True)
class LogUpGkrOutput:
    """Final numerator/denominator MLEs after all transitions."""

    numerator: Array
    denominator: Array

JaggedGkrLayer dataclass

One jagged fractional-sum layer at a capacity width, stored batch-major.

The four MLEs are flat over a static capacity width >= sum(row_counts): the live prefix holds all rows of batch element 0, then batch element 1, and so on; every slot past sum(row_counts) is DEAD and laid zero. The counts ride as one traced i32[num_batches] vector, so no per-input layout value keys a compile — transitions derive their gathers in-trace and the layer rounds take the schedule as traced operands, keying on the capacity shapes alone. An exactly-sized layer (width == sum(row_counts)) is just the zero-slack case; the batch count must be a power of two (the consumer pads its interaction list).

The dead tail is never read: a transition resolves every non-live source through its gather sentinel, and the layer rounds bound their reads by the live-pair operand. Zero (not the neutral fraction) so a reduction that sweeps a fixed-width buffer picks up nothing from the dead region. Truncation-safety is likewise the consumer's obligation — a host guard cannot read the traced counts, so a transition schedule must dominate ceil(rc / 2) pointwise and its width must hold its count sum for every input the consumer admits.

Source code in zorch/logup_gkr/circuit.py

@partial(
    frx.tree_util.register_dataclass,
    data_fields=[
        "numerator_0",
        "numerator_1",
        "denominator_0",
        "denominator_1",
        "row_counts",
    ],
    meta_fields=[],
)
@dataclass(frozen=True)
class JaggedGkrLayer:
    """One jagged fractional-sum layer at a capacity width, stored batch-major.

    The four MLEs are flat over a static capacity `width >= sum(row_counts)`:
    the live prefix holds all rows of batch element 0, then batch element 1,
    and so on; every slot past `sum(row_counts)` is DEAD and laid zero. The
    counts ride as one traced i32[num_batches] vector, so no per-input layout
    value keys a compile — transitions derive their gathers in-trace and the
    layer rounds take the schedule as traced operands, keying on the capacity
    shapes alone. An exactly-sized layer (`width == sum(row_counts)`) is just
    the zero-slack case; the batch count must be a power of two (the consumer
    pads its interaction list).

    The dead tail is never read: a transition resolves every non-live source
    through its gather sentinel, and the layer rounds bound their reads by
    the live-pair operand. Zero (not the neutral fraction) so a reduction
    that sweeps a fixed-width buffer picks up nothing from the dead region.
    Truncation-safety is likewise the consumer's obligation — a host guard
    cannot read the traced counts, so a transition schedule must dominate
    `ceil(rc / 2)` pointwise and its width must hold its count sum for every
    input the consumer admits.
    """

    numerator_0: Array
    numerator_1: Array
    denominator_0: Array
    denominator_1: Array
    row_counts: Array

    def __post_init__(self) -> None:
        # Shape-only checks: `register_dataclass` reruns this during
        # unflatten, so no value may be branched on here — and the leaves are
        # not always arrays: AOT lowering (`jit(f).lower(layer)`) rebuilds the
        # tree with `frx.stages.ArgInfo` leaves, which expose `shape`/`dtype`
        # but no `ndim`. Stick to `.shape`.
        log2_strict_usize(self.num_batches)
        if len(self.row_counts.shape) != 1:
            raise ValueError(
                f"row_counts must be a flat vector, got {self.row_counts.shape}"
            )
        width = self.width
        for name in (
            "numerator_0",
            "numerator_1",
            "denominator_0",
            "denominator_1",
        ):
            shape = getattr(self, name).shape
            if shape != (width,):
                raise ValueError(
                    f"the four MLEs must share one capacity width; "
                    f"{name} is {shape}, expected ({width},)"
                )

    @property
    def num_batches(self) -> int:
        return self.row_counts.shape[0]

    @property
    def num_batch_variables(self) -> int:
        return log2_strict_usize(self.num_batches)

    @property
    def width(self) -> int:
        return self.numerator_0.shape[0]

layer_transition

layer_transition(layer: GkrLayer) -> GkrLayer

Fold one row variable: sum each LSB-adjacent fraction pair (stride 2).

Requires a row variable to fold.

Source code in zorch/logup_gkr/circuit.py

def layer_transition(layer: GkrLayer) -> GkrLayer:
    """Fold one row variable: sum each LSB-adjacent fraction pair (stride 2).

    Requires a row variable to fold.
    """
    if layer.num_row_variables < 1:
        raise ValueError("no row variable left to fold")
    rn0, rn1, rd0, rd1 = _fold_pairs(
        layer.numerator_0,
        layer.numerator_1,
        layer.denominator_0,
        layer.denominator_1,
    )
    return GkrLayer(
        numerator_0=rn0,
        numerator_1=rn1,
        denominator_0=rd0,
        denominator_1=rd1,
        num_batch_variables=layer.num_batch_variables,
    )

build_pyramid

build_pyramid(first: GkrLayer) -> list[GkrLayer]

Eager fold from first (most row variables) down to the batch floor.

Eager rather than one fused program: the full pyramid does not fit one @jit at scale, and each transition's output feeds the next layer's per-variable sumcheck independently.

Source code in zorch/logup_gkr/circuit.py

def build_pyramid(first: GkrLayer) -> list[GkrLayer]:
    """Eager fold from `first` (most row variables) down to the batch floor.

    Eager rather than one fused program: the full pyramid does not fit one
    `@jit` at scale, and each transition's output feeds the next layer's
    per-variable sumcheck independently.
    """
    layers = [first]
    while layers[-1].num_row_variables > 0:
        layers.append(layer_transition(layers[-1]))
    return layers

extract_outputs

extract_outputs(layer: GkrLayer) -> LogUpGkrOutput

Interleave the two children of the floor layer into the output MLEs.

At num_row_variables == 0 each index is one batch element's fraction with a 0-child and a 1-child; interleaving recovers the MLE over the batch variables plus one.

Source code in zorch/logup_gkr/circuit.py

def extract_outputs(layer: GkrLayer) -> LogUpGkrOutput:
    """Interleave the two children of the floor layer into the output MLEs.

    At `num_row_variables == 0` each index is one batch element's fraction with a
    0-child and a 1-child; interleaving recovers the MLE over the batch
    variables plus one.
    """
    if layer.num_row_variables != 0:
        raise ValueError(
            f"extract_outputs expects the batch floor (num_row_variables "
            f"== 0), got {layer.num_row_variables}"
        )
    return LogUpGkrOutput(
        numerator=_interleave(layer.numerator_0, layer.numerator_1),
        denominator=_interleave(layer.denominator_0, layer.denominator_1),
    )

jagged_layer_transition

jagged_layer_transition(
    layer: JaggedGkrLayer,
    out_row_counts: Array | Sequence[int],
    out_width: int | None = None,
) -> JaggedGkrLayer

Fold one row variable per segment into a fresh out_width capacity.

Odd segments pre-pad with the additive-identity fraction (n=0, d=1) so the stride-2 fold never pairs across a batch boundary; slots between a segment's folded count and its out count are live neutral padding; the dead region past sum(out_row_counts) lands zero.

out_row_counts is the consumer's halving policy: a host sequence for a statically-known schedule (out_width defaults to its sum — the zero-slack layout), or the policy evaluated on ITS traced counts, with out_width the static capacity holding it. A schedule that truncates a segment's folded rows cannot be rejected host-side (the counts are traced), so the consumer owns that guarantee — its policy must dominate ceil(rc / 2) pointwise for every input it admits.

Source code in zorch/logup_gkr/circuit.py

def jagged_layer_transition(
    layer: JaggedGkrLayer,
    out_row_counts: Array | Sequence[int],
    out_width: int | None = None,
) -> JaggedGkrLayer:
    """Fold one row variable per segment into a fresh `out_width` capacity.

    Odd segments pre-pad with the additive-identity fraction (n=0, d=1) so
    the stride-2 fold never pairs across a batch boundary; slots between a
    segment's folded count and its out count are live neutral padding; the
    dead region past `sum(out_row_counts)` lands zero.

    `out_row_counts` is the consumer's halving policy: a host sequence for a
    statically-known schedule (`out_width` defaults to its sum — the
    zero-slack layout), or the policy evaluated on ITS traced counts, with
    `out_width` the static capacity holding it. A schedule that truncates a
    segment's folded rows cannot be rejected host-side (the counts are
    traced), so the consumer owns that guarantee — its policy must dominate
    `ceil(rc / 2)` pointwise for every input it admits.
    """
    if isinstance(out_row_counts, Array):
        if out_width is None:
            raise ValueError("a traced schedule needs an explicit out_width capacity")
        counts = out_row_counts
        num_out = counts.shape[0]
    else:
        host = tuple(int(rc) for rc in out_row_counts)
        counts = _counts_operand(host)
        num_out = len(host)
        if out_width is None:
            out_width = sum(host)
    if num_out != layer.num_batches:
        raise ValueError(
            f"schedule must cover all {layer.num_batches} batches, got "
            f"{num_out} entries"
        )
    rn0, rn1, rd0, rd1 = _jagged_transition_core(
        layer.numerator_0,
        layer.numerator_1,
        layer.denominator_0,
        layer.denominator_1,
        layer.row_counts,
        counts,
        out_width=out_width,
    )
    return JaggedGkrLayer(
        numerator_0=rn0,
        numerator_1=rn1,
        denominator_0=rd0,
        denominator_1=rd1,
        row_counts=counts,
    )

build_jagged_pyramid

build_jagged_pyramid(
    first: JaggedGkrLayer,
    schedules: Sequence[tuple[Array, int] | Sequence[int]],
) -> list[JaggedGkrLayer]

Build the jagged pyramid [first, ..., floor], folding one row variable per transition. schedules[k] is transition k's policy — (out_row_counts, out_width) for a traced schedule, or a bare host sequence at its zero-slack width — the same argument jagged_layer_transition takes. Each transition is its own dispatch, so the layers land as separate per-layer buffers. Peak residency is ~2H -- every natural-width layer is a required GKR input, live until its top-down sumcheck -- split across ~depth buffers a pooling allocator can seat individually; a single contiguous 2H alloc exceeds what BFC can place on wide shards (#468). One compile per distinct (in_width, out_width) pair, shared by every input of the capacity class.

Source code in zorch/logup_gkr/circuit.py

def build_jagged_pyramid(
    first: JaggedGkrLayer,
    schedules: Sequence[tuple[Array, int] | Sequence[int]],
) -> list[JaggedGkrLayer]:
    """Build the jagged pyramid `[first, ..., floor]`, folding one row
    variable per transition. `schedules[k]` is transition `k`'s policy —
    `(out_row_counts, out_width)` for a traced schedule, or a bare host
    sequence at its zero-slack width — the same argument
    `jagged_layer_transition` takes. Each transition is its own dispatch, so
    the layers land as separate per-layer buffers. Peak residency is ~2*H --
    every natural-width layer is a required GKR input, live until its
    top-down sumcheck -- split across ~depth buffers a pooling allocator can
    seat individually; a single contiguous 2*H alloc exceeds what BFC can
    place on wide shards (#468). One compile per distinct
    (in_width, out_width) pair, shared by every input of the capacity class.
    """
    layers = [first]
    layer = first
    for schedule in schedules:
        if (
            isinstance(schedule, tuple)
            and len(schedule) == 2
            and isinstance(schedule[0], Array)
        ):
            layer = jagged_layer_transition(layer, schedule[0], schedule[1])
        else:
            layer = jagged_layer_transition(layer, schedule)
        layers.append(layer)
    return layers

extract_jagged_outputs

extract_jagged_outputs(
    layer: JaggedGkrLayer,
) -> LogUpGkrOutput

Interleave the floor layer's children into the output MLEs.

The floor is row counts all 1 -- one fraction pair per batch element -- the jagged dual of extract_outputs's num_row_variables == 0 precondition. Row counts are traced, so the gate is the static dual the layout implies: a fully-live all-ones layer is exactly width == num_batches (any live count above 1 would need more width; a dead slot would mean a count of 0, which no saturating fold produces). A schedule that stops higher folds the rest down with jagged_layer_transition first; how far to fold is the consumer's call.

Source code in zorch/logup_gkr/circuit.py

def extract_jagged_outputs(layer: JaggedGkrLayer) -> LogUpGkrOutput:
    """Interleave the floor layer's children into the output MLEs.

    The floor is row counts all 1 -- one fraction pair per batch element --
    the jagged dual of `extract_outputs`'s `num_row_variables == 0`
    precondition. Row counts are traced, so the gate is the static dual the
    layout implies: a fully-live all-ones layer is exactly `width ==
    num_batches` (any live count above 1 would need more width; a dead slot
    would mean a count of 0, which no saturating fold produces). A schedule
    that stops higher folds the rest down with `jagged_layer_transition`
    first; how far to fold is the consumer's call.
    """
    if layer.width != layer.num_batches:
        raise ValueError(
            f"extract_jagged_outputs expects the batch floor (width == "
            f"num_batches == {layer.num_batches}), got width {layer.width}"
        )
    return LogUpGkrOutput(
        numerator=_interleave(layer.numerator_0, layer.numerator_1),
        denominator=_interleave(layer.denominator_0, layer.denominator_1),
    )