zorch.coding.reed_solomon

Reed-Solomon as a FoldableCode: low-degree extension via the native NTT.

encode reads the message as the message_len low-order coefficients of a polynomial, zero-pads to block_len, and evaluates it on the order-block_len two-adic subgroup (or a coset of it). The evaluation is frx.lax.ntt — the XLA-native NTT — which lowers to one fused kernel and auto-decomposes extension fields into prime-field NTTs.

There is deliberately no hand-rolled butterfly: a fnp butterfly would be log(n) unfused kernels the compiler cannot recognize as an NTT. Reed-Solomon hands its evaluation to the native op, the way poseidon2 hands its algebra to XLA rather than fusing it by pattern-match.

fri_fold is the codeword fold shared by every FRI-style scheme (FRI, Basefold, WHIR, STARK); the fold half of the seam delegates to it. It lives in this module so the fold's x-coordinates stay the same evaluation domain the encoder used. The arbitrary-fold-factor (k-ary) generalization is fri_fold_k plus the KFoldableCode group seam — additive to the binary conjugate-pair fold, which stays the closed-form butterfly (see coding.md).

ReedSolomon

Reed-Solomon code over dtype; implements FoldableCode.

block_len = message_len * blowup (both powers of two). With coset_shift set to a field element outside the subgroup, the codeword is the message polynomial evaluated on the coset coset_shift * <subgroup> rather than the subgroup itself — FRI/STARK want an evaluation domain disjoint from the trace domain. The shift is supplied by the caller, so the code carries no field-generator table.

Source code in zorch/coding/reed_solomon.py

class ReedSolomon:
    """Reed-Solomon code over `dtype`; implements FoldableCode.

    `block_len = message_len * blowup` (both powers of two). With `coset_shift`
    set to a field element outside the subgroup, the codeword is the message
    polynomial evaluated on the coset `coset_shift * <subgroup>` rather than the
    subgroup itself — FRI/STARK want an evaluation domain disjoint from the
    trace domain. The shift is supplied by the caller, so the code carries no
    field-generator table.
    """

    def __init__(
        self,
        message_len: int,
        blowup: int,
        dtype: Any,
        *,
        coset_shift: Array | None = None,
        fold_factor: int = 2,
        generator: int | None = None,
    ) -> None:
        if not is_power_of_two(message_len):
            raise ValueError(f"message_len must be a power of two, got {message_len}")
        if not is_power_of_two(blowup):
            raise ValueError(f"blowup must be a power of two, got {blowup}")
        if not is_power_of_two(fold_factor) or fold_factor < 2:
            raise ValueError(
                f"fold_factor must be a power of two >= 2, got {fold_factor}"
            )
        self.message_len = message_len
        self.block_len = message_len * blowup
        self.dtype = dtype
        self.coset_shift = coset_shift
        # Subgroup generator selecting the NTT root (`gen^((p-1)/n)`); None uses
        # the dtype's canonical root. Lets a consumer match an external prover's
        # domain order without re-permuting the codeword. Honoured on the
        # encode/LDE surface (`encode`, `domain`); the FRI fold path
        # (`fold`/`fold_values`) is canonical-order only — generator there is a
        # follow-up.
        self.generator = generator
        # The k-ary fold's static factor (KFoldableCode). The binary pair seam
        # ignores it; default 2 keeps the binary path's jit-zone key unchanged.
        self.fold_factor = fold_factor
        self._log2_fold_factor = log2_strict_usize(fold_factor)
        # Coset eval scales coeffs by [1, h, h^2, ..., h^{n-1}]; precompute it
        # once since h and n are fixed (paid eagerly at construction). `powers`
        # log-doubles rather than `fnp.cumprod`/`fnp.arange` (iota raises on
        # extension dtypes).
        self._coset_powers = None
        if coset_shift is not None:
            self._coset_powers = powers(fnp.asarray(coset_shift, dtype), self.block_len)
        self._key: tuple | None = None
        # A binary field has characteristic 2 — its multiplicative group has odd
        # order and no 2^m-th roots of unity, so `lax.ntt` runs the LCH *additive*
        # NTT (novel basis over an F2-affine subspace) instead of the
        # multiplicative one. `encode` is unchanged (`lax.ntt` dispatches by
        # dtype), but `eval_point` needs the additive tensor — its factors shift
        # from the geometric `(dₛ^{2^i})` to the subspace polynomials `(Ŵ_i(dₛ))`
        # — and the multiplicative fold/coset do not apply.
        self._binary = is_binary_field(dtype)
        if self._binary and coset_shift is not None:
            raise NotImplementedError(
                "coset Reed-Solomon over a binary field (an additive coset) is "
                "unimplemented; only the base F2-subspace is supported"
            )
        # The additive eval-point tensor is the k basis-power codewords `Ŵ_i`
        # over the domain (the geometric squaring the prime path uses does not
        # apply in the novel basis). Built lazily on first `eval_point`:
        # encode-only consumers (Ligero matrix commits, plain LDE) never pay
        # the k extra encodes, and construction stays cheap enough to build
        # codes per level (the #214 value-keyed-instances contract). Prime
        # fields keep the on-the-fly geometric form and never build a table.
        self._binary_eval_table: Array | None = None

    # Value equality/hash for static jit-zone keys — the LinearCode seam
    # contract (#214). The key is cached host-side because jit dispatch
    # compares static args per call, and `tobytes` on the live coset-shift
    # array would cost a device->host sync each time (the Poseidon2Params
    # pattern).
    def _value_key(self) -> tuple:
        if self._key is None:
            shift = (
                None
                if self.coset_shift is None
                else np.asarray(self.coset_shift).tobytes()
            )
            self._key = (
                self.message_len,
                self.block_len,
                self.dtype,
                shift,
                self.fold_factor,
                self.generator,
            )
        return self._key

    def __eq__(self, other: object) -> bool:
        if self is other:
            return True
        if not isinstance(other, ReedSolomon):
            return NotImplemented
        return self._value_key() == other._value_key()

    def __hash__(self) -> int:
        return hash(self._value_key())

    def encode(self, message: Array) -> Array:
        if message.shape[-1] != self.message_len:
            raise ValueError(
                f"message last axis must be {self.message_len}, "
                f"got {message.shape[-1]}"
            )
        n = self.block_len
        tail = message.shape[:-1] + (n - self.message_len,)
        coeffs = fnp.concatenate([message, fnp.zeros(tail, self.dtype)], axis=-1)
        if self._coset_powers is not None:
            coeffs = coeffs * self._coset_powers
        return lax.ntt(coeffs, ntt_type="NTT", ntt_length=n, generator=self.generator)

    def extend(self, evals: Array) -> Array:
        """LDE `evals` — the `message_len` evaluations on the base subgroup — to
        the `block_len` coset codeword, transforming the last axis (any leading
        batch axes are preserved). Natural order in and out.

        Interpolates to coefficients (`intt`), then coset-evaluates (`encode`) —
        i.e. `encode(intt(evals))`.
        """
        if evals.shape[-1] != self.message_len:
            raise ValueError(
                f"evals last axis must be {self.message_len}, got {evals.shape[-1]}"
            )
        coeffs = lax.ntt(
            evals,
            ntt_type="INTT",
            ntt_length=self.message_len,
            generator=self.generator,
        )
        return self.encode(coeffs)

    def domain(self) -> Array:
        """The points `encode` evaluates on, coset shift included."""
        return eval_domain(
            self.dtype,
            self.block_len,
            shift=self.coset_shift,
            generator=self.generator,
        )

    def eval_point(self, positions: Array) -> Array:
        """TensorCode seam: the multilinear point `p_s` that codeword coordinate
        `positions` evaluates, so that
        `encode(w)[s] == eval_mle(mle_coeffs_to_evals(w), eval_point(s))` — which
        turns the Ligerito proximity RHS into a point-eval of the committed `w`.

        For the multiplicative NTT `encode` is monomial-basis
        (`encode(w)[s] = Σ_j w[j]·dₛʲ`, `dₛ = domain()[s]`), so the generator row
        `(1, dₛ, dₛ², …)` factors as the geometric tensor
        `p_s = (dₛ^{2^{k-1}}, …, dₛ², dₛ)` — MSB-first to match `eval_mle`'s
        lexicographic eq order. The `dₛ^{2^i}` are built by repeated squaring so
        no array exponent is taken (field dtypes reject `fnp.power`).

        For a binary field `encode` is the additive NTT (novel basis), so the
        tensor factors are the subspace polynomials `Ŵ_i` instead; those are the
        basis-power codewords, gathered from a table built on first use."""
        if self._binary:
            if self._binary_eval_table is None:
                # Concrete even under an outer trace: a traced table stored on
                # self would leak the tracer into later calls.
                with frx.ensure_compile_time_eval():
                    self._binary_eval_table = self._build_binary_eval_table()
            return self._binary_eval_table[positions]  # (*positions.shape, k)
        k = log2_strict_usize(self.message_len)
        cur = self.domain()[positions]  # (*positions.shape,) evaluation point(s)
        cols = []
        for _ in range(k):
            cols.append(cur)
            cur = cur * cur
        # MSB-first: variable 0 binds the highest power dₛ^{2^{k-1}}.
        return fnp.stack(cols[::-1], axis=-1)  # (*positions.shape, k)

    def _build_binary_eval_table(self) -> Array:
        """The `(block_len, k)` additive eval-point table, column `i` = the
        subspace polynomial `Ŵ_{k-1-i}` over the domain (MSB-first).

        `Ŵ_i(dₛ) == encode(e_{2^i})[s]`: the additive NTT of the unit novel-basis
        coefficient `e_{2^i}` is that subspace polynomial over the whole domain,
        so the tensor factors are read off the encoder itself. A per-query
        subset-XOR reconstruction (each `Ŵ_i` is F2-linear) would be lighter but
        needs an int→field select, which the binary-field lowering does not yet
        support; the resident gather lowers cleanly. Built from concrete arrays,
        so no `.at[].set` (scatter is unsupported here)."""
        k = log2_strict_usize(self.message_len)
        units = fnp.asarray(
            [
                [1 if c == (1 << i) else 0 for c in range(self.message_len)]
                for i in range(k)
            ],
            dtype=self.dtype,
        )
        basis_codewords = self.encode(units)  # (k, block_len) = Ŵ_i over the domain
        return basis_codewords[::-1].T  # (block_len, k), MSB-first

    def fold(self, codeword: Array, beta: Array) -> Array:
        """FoldableCode fold: natural-order `(x, -x)` conjugate pairs. The layer
        level — and with it the coset shift — is read off the codeword length."""
        self._reject_binary_fold()
        level = log2_strict_usize(self.block_len // codeword.shape[0])
        return fri_fold(codeword, beta, shift=self._level_shift(level))

    def _reject_binary_fold(self) -> None:
        """The FRI fold is the multiplicative `(x, -x)` conjugate pair; the
        additive-NTT analog is a different (subspace) fold and is unimplemented.
        A binary-field Reed-Solomon is used as a `TensorCode` (Ligerito, which
        re-commits per level), not a `FoldableCode` — guard rather than fold
        silently over the wrong domain."""
        if self._binary:
            raise NotImplementedError(
                "additive-NTT (binary-field) fold is unimplemented; a binary-field "
                "Reed-Solomon is a TensorCode (Ligerito), not a FoldableCode"
            )

    def fold_values(
        self, lo: Array, hi: Array, beta: Array, positions: Array, level: int
    ) -> Array:
        """Fold opened pairs of layer `level`; the x-coordinates are the first
        half of the layer's (level-times-squared) evaluation domain."""
        self._reject_binary_fold()
        domain = eval_domain(
            self.dtype, self.block_len >> level, shift=self._level_shift(level)
        )
        return fri_fold_values(lo, hi, beta, domain[positions])

    def pair_leaves(self, codeword: Array) -> Array:
        """Natural order: conjugates sit a half-layer apart, so leaf `p` is
        `(codeword[p], codeword[p + half])`."""
        half = codeword.shape[0] // 2
        return fnp.stack([codeword[:half], codeword[half:]], axis=1)

    def check_final(self, final: Array, claim: Array) -> Array:
        """A message-length-1 RS codeword is the constant polynomial on any
        domain, so base-code membership and message == `claim` collapse into one
        comparison."""
        return fnp.all(final == claim)

    def pair_indices(self, positions: Array, level: int) -> tuple[Array, Array]:
        """Natural order: the conjugates of layer `level` sit a half-layer
        apart, and the lo index is the landing index itself."""
        return positions, positions + (self.block_len >> (level + 1))

    def layer_positions(self, positions: Array, num_rounds: int) -> list[Array]:
        """Natural order: `a_i = q_i mod (n / 2^{i+1})` with `q_0 = positions`,
        `q_{i+1} = a_i`, elementwise over the query axis."""
        indices = []
        q = positions
        for i in range(num_rounds):
            a = q % (self.block_len >> (i + 1))
            indices.append(a)
            q = a
        return indices

    def _level_shift(self, level: int) -> Array | None:
        """Layer `level`'s domain shift, `coset_shift^(2^level)` — each fold
        lands on the squared domain, squaring the shift with it."""
        if self.coset_shift is None:
            return None
        shift = self.coset_shift
        for _ in range(level):
            shift = shift * shift
        return shift

    # --- KFoldableCode: the k-ary fold seam, additive to the binary pair seam
    # above. A k-ary fold groups the k entries of one folded point's k-th-root
    # coset {p, p + n/k, ..., p + (k-1)n/k} (the natural-order generalization of
    # the (x, -x) conjugate pair) and Lagrange-interpolates them at beta via
    # `fri_fold_k`'s `points=` form; the binary pair keeps its cheaper closed-form
    # butterfly.

    def fold_group(self, codeword: Array, beta: Array) -> Array:
        """KFoldableCode fold: regroup the layer into k-th-root cosets and
        Lagrange-fold each at `beta`, dividing the length by `fold_factor`. The
        level — and with it the coset shift — is read off the codeword length."""
        k = self.fold_factor
        n = codeword.shape[0]
        # Fail loud at the seam boundary, the k-ary twin of fri_fold's n<2 guard:
        # a layer too short to form one k-group would otherwise die in
        # `_regroup`'s reshape with an opaque shape error. A real codeword length
        # is a power of two, so n>=k already implies k divides n; n is static
        # (shape), so this is vmap-safe.
        if n < k:
            raise ValueError(
                f"fold_group requires a codeword length >= fold_factor {k}, got {n}"
            )
        level = log2_strict_usize(self.block_len // n) // self._log2_fold_factor
        domain = eval_domain(
            _base_dtype(codeword.dtype), n, shift=self._group_level_shift(level)
        )
        return self._fold_groups(
            self._regroup(codeword, k), self._regroup(domain, k), beta
        )

    def group_leaves(self, codeword: Array) -> Array:
        """Natural order: a k-th-root coset sits a sub-layer (`n // k`) apart, so
        leaf `p` is `(codeword[p], codeword[p + n/k], ..., codeword[p +
        (k-1)n/k])` — the k-ary `pair_leaves`."""
        return self._regroup(codeword, self.fold_factor)

    def group_indices(self, positions: Array, level: int) -> tuple[Array, ...]:
        """Natural order: the k-th-root coset of layer `level` whose fold lands
        at `positions` sits a sub-layer apart, and `positions` itself is the
        landing (first) index."""
        sub = self.block_len >> (self._log2_fold_factor * (level + 1))
        return tuple(positions + m * sub for m in range(self.fold_factor))

    def fold_group_values(
        self, group: Array, beta: Array, positions: Array, level: int
    ) -> Array:
        """Fold opened k-groups of layer `level`; the x-coordinates are the
        group's points on the layer's `(level`-times-`k`-th-powered) domain."""
        n = self.block_len >> (self._log2_fold_factor * level)
        domain = eval_domain(
            _base_dtype(self.dtype), n, shift=self._group_level_shift(level)
        )
        points = domain[fnp.stack(self.group_indices(positions, level), axis=-1)]
        return self._fold_groups(group, points, beta)

    def group_layer_positions(self, positions: Array, num_rounds: int) -> list[Array]:
        """Natural order: `a_i = q_i mod (n / k^{i+1})` with `q_0 = positions`,
        `q_{i+1} = a_i`, elementwise over the query axis — the k-ary
        `layer_positions`."""
        indices = []
        q = positions
        for i in range(num_rounds):
            a = q % (self.block_len >> (self._log2_fold_factor * (i + 1)))
            indices.append(a)
            q = a
        return indices

    def _fold_groups(self, groups: Array, points: Array, beta: Array) -> Array:
        """vmap the single-group Lagrange fold (`fri_fold_k` with `points`) over
        the group/query axis — the one place fold_group and fold_group_values
        share, each supplying its own `groups`/`points` (full layer vs opened
        queries)."""
        return frx.vmap(lambda g, p: fri_fold_k(g, beta, points=p))(groups, points)

    def _regroup(self, layer: Array, k: int) -> Array:
        """Reshape a length-`n` layer into its `[n // k, k]` k-th-root cosets:
        row `p` is `layer[[p, p + n/k, ..., p + (k-1)n/k]]`. `reshape(k, n//k).T`
        gathers that coset without an index array."""
        return layer.reshape(k, layer.shape[0] // k).T

    def _group_level_shift(self, level: int) -> Array | None:
        """Layer `level`'s domain shift under k-ary folding, `coset_shift^(k^level)`
        — each fold lands on the `k`-th-powered domain, raising the shift with it.
        Since `k^level = 2^(log2(k)·level)`, this is exactly the binary
        `_level_shift` at `log2(k)·level` squarings, so it delegates there."""
        return self._level_shift(self._log2_fold_factor * level)

extend

extend(evals: Array) -> Array

LDE evals — the message_len evaluations on the base subgroup — to the block_len coset codeword, transforming the last axis (any leading batch axes are preserved). Natural order in and out.

Interpolates to coefficients (intt), then coset-evaluates (encode) — i.e. encode(intt(evals)).

Source code in zorch/coding/reed_solomon.py

def extend(self, evals: Array) -> Array:
    """LDE `evals` — the `message_len` evaluations on the base subgroup — to
    the `block_len` coset codeword, transforming the last axis (any leading
    batch axes are preserved). Natural order in and out.

    Interpolates to coefficients (`intt`), then coset-evaluates (`encode`) —
    i.e. `encode(intt(evals))`.
    """
    if evals.shape[-1] != self.message_len:
        raise ValueError(
            f"evals last axis must be {self.message_len}, got {evals.shape[-1]}"
        )
    coeffs = lax.ntt(
        evals,
        ntt_type="INTT",
        ntt_length=self.message_len,
        generator=self.generator,
    )
    return self.encode(coeffs)

domain

domain() -> Array

The points encode evaluates on, coset shift included.

Source code in zorch/coding/reed_solomon.py

def domain(self) -> Array:
    """The points `encode` evaluates on, coset shift included."""
    return eval_domain(
        self.dtype,
        self.block_len,
        shift=self.coset_shift,
        generator=self.generator,
    )

eval_point

eval_point(positions: Array) -> Array

TensorCode seam: the multilinear point p_s that codeword coordinate positions evaluates, so that encode(w)[s] == eval_mle(mle_coeffs_to_evals(w), eval_point(s)) — which turns the Ligerito proximity RHS into a point-eval of the committed w.

For the multiplicative NTT encode is monomial-basis (encode(w)[s] = Σ_j w[j]·dₛʲ, dₛ = domain()[s]), so the generator row (1, dₛ, dₛ², …) factors as the geometric tensor p_s = (dₛ^{2^{k-1}}, …, dₛ², dₛ) — MSB-first to match eval_mle's lexicographic eq order. The dₛ^{2^i} are built by repeated squaring so no array exponent is taken (field dtypes reject fnp.power).

For a binary field encode is the additive NTT (novel basis), so the tensor factors are the subspace polynomials Ŵ_i instead; those are the basis-power codewords, gathered from a table built on first use.

Source code in zorch/coding/reed_solomon.py

def eval_point(self, positions: Array) -> Array:
    """TensorCode seam: the multilinear point `p_s` that codeword coordinate
    `positions` evaluates, so that
    `encode(w)[s] == eval_mle(mle_coeffs_to_evals(w), eval_point(s))` — which
    turns the Ligerito proximity RHS into a point-eval of the committed `w`.

    For the multiplicative NTT `encode` is monomial-basis
    (`encode(w)[s] = Σ_j w[j]·dₛʲ`, `dₛ = domain()[s]`), so the generator row
    `(1, dₛ, dₛ², …)` factors as the geometric tensor
    `p_s = (dₛ^{2^{k-1}}, …, dₛ², dₛ)` — MSB-first to match `eval_mle`'s
    lexicographic eq order. The `dₛ^{2^i}` are built by repeated squaring so
    no array exponent is taken (field dtypes reject `fnp.power`).

    For a binary field `encode` is the additive NTT (novel basis), so the
    tensor factors are the subspace polynomials `Ŵ_i` instead; those are the
    basis-power codewords, gathered from a table built on first use."""
    if self._binary:
        if self._binary_eval_table is None:
            # Concrete even under an outer trace: a traced table stored on
            # self would leak the tracer into later calls.
            with frx.ensure_compile_time_eval():
                self._binary_eval_table = self._build_binary_eval_table()
        return self._binary_eval_table[positions]  # (*positions.shape, k)
    k = log2_strict_usize(self.message_len)
    cur = self.domain()[positions]  # (*positions.shape,) evaluation point(s)
    cols = []
    for _ in range(k):
        cols.append(cur)
        cur = cur * cur
    # MSB-first: variable 0 binds the highest power dₛ^{2^{k-1}}.
    return fnp.stack(cols[::-1], axis=-1)  # (*positions.shape, k)

fold

fold(codeword: Array, beta: Array) -> Array

FoldableCode fold: natural-order (x, -x) conjugate pairs. The layer level — and with it the coset shift — is read off the codeword length.

Source code in zorch/coding/reed_solomon.py

def fold(self, codeword: Array, beta: Array) -> Array:
    """FoldableCode fold: natural-order `(x, -x)` conjugate pairs. The layer
    level — and with it the coset shift — is read off the codeword length."""
    self._reject_binary_fold()
    level = log2_strict_usize(self.block_len // codeword.shape[0])
    return fri_fold(codeword, beta, shift=self._level_shift(level))

fold_values

fold_values(
    lo: Array,
    hi: Array,
    beta: Array,
    positions: Array,
    level: int,
) -> Array

Fold opened pairs of layer level; the x-coordinates are the first half of the layer's (level-times-squared) evaluation domain.

Source code in zorch/coding/reed_solomon.py

def fold_values(
    self, lo: Array, hi: Array, beta: Array, positions: Array, level: int
) -> Array:
    """Fold opened pairs of layer `level`; the x-coordinates are the first
    half of the layer's (level-times-squared) evaluation domain."""
    self._reject_binary_fold()
    domain = eval_domain(
        self.dtype, self.block_len >> level, shift=self._level_shift(level)
    )
    return fri_fold_values(lo, hi, beta, domain[positions])

pair_leaves

pair_leaves(codeword: Array) -> Array

Natural order: conjugates sit a half-layer apart, so leaf p is (codeword[p], codeword[p + half]).

Source code in zorch/coding/reed_solomon.py

def pair_leaves(self, codeword: Array) -> Array:
    """Natural order: conjugates sit a half-layer apart, so leaf `p` is
    `(codeword[p], codeword[p + half])`."""
    half = codeword.shape[0] // 2
    return fnp.stack([codeword[:half], codeword[half:]], axis=1)

check_final

check_final(final: Array, claim: Array) -> Array

A message-length-1 RS codeword is the constant polynomial on any domain, so base-code membership and message == claim collapse into one comparison.

Source code in zorch/coding/reed_solomon.py

def check_final(self, final: Array, claim: Array) -> Array:
    """A message-length-1 RS codeword is the constant polynomial on any
    domain, so base-code membership and message == `claim` collapse into one
    comparison."""
    return fnp.all(final == claim)

pair_indices

pair_indices(
    positions: Array, level: int
) -> tuple[Array, Array]

Natural order: the conjugates of layer level sit a half-layer apart, and the lo index is the landing index itself.

Source code in zorch/coding/reed_solomon.py

def pair_indices(self, positions: Array, level: int) -> tuple[Array, Array]:
    """Natural order: the conjugates of layer `level` sit a half-layer
    apart, and the lo index is the landing index itself."""
    return positions, positions + (self.block_len >> (level + 1))

layer_positions

layer_positions(
    positions: Array, num_rounds: int
) -> list[Array]

Natural order: a_i = q_i mod (n / 2^{i+1}) with q_0 = positions, q_{i+1} = a_i, elementwise over the query axis.

Source code in zorch/coding/reed_solomon.py

def layer_positions(self, positions: Array, num_rounds: int) -> list[Array]:
    """Natural order: `a_i = q_i mod (n / 2^{i+1})` with `q_0 = positions`,
    `q_{i+1} = a_i`, elementwise over the query axis."""
    indices = []
    q = positions
    for i in range(num_rounds):
        a = q % (self.block_len >> (i + 1))
        indices.append(a)
        q = a
    return indices

fold_group

fold_group(codeword: Array, beta: Array) -> Array

KFoldableCode fold: regroup the layer into k-th-root cosets and Lagrange-fold each at beta, dividing the length by fold_factor. The level — and with it the coset shift — is read off the codeword length.

Source code in zorch/coding/reed_solomon.py

def fold_group(self, codeword: Array, beta: Array) -> Array:
    """KFoldableCode fold: regroup the layer into k-th-root cosets and
    Lagrange-fold each at `beta`, dividing the length by `fold_factor`. The
    level — and with it the coset shift — is read off the codeword length."""
    k = self.fold_factor
    n = codeword.shape[0]
    # Fail loud at the seam boundary, the k-ary twin of fri_fold's n<2 guard:
    # a layer too short to form one k-group would otherwise die in
    # `_regroup`'s reshape with an opaque shape error. A real codeword length
    # is a power of two, so n>=k already implies k divides n; n is static
    # (shape), so this is vmap-safe.
    if n < k:
        raise ValueError(
            f"fold_group requires a codeword length >= fold_factor {k}, got {n}"
        )
    level = log2_strict_usize(self.block_len // n) // self._log2_fold_factor
    domain = eval_domain(
        _base_dtype(codeword.dtype), n, shift=self._group_level_shift(level)
    )
    return self._fold_groups(
        self._regroup(codeword, k), self._regroup(domain, k), beta
    )

group_leaves

group_leaves(codeword: Array) -> Array

Natural order: a k-th-root coset sits a sub-layer (n // k) apart, so leaf p is (codeword[p], codeword[p + n/k], ..., codeword[p + (k-1)n/k]) — the k-ary pair_leaves.

Source code in zorch/coding/reed_solomon.py

def group_leaves(self, codeword: Array) -> Array:
    """Natural order: a k-th-root coset sits a sub-layer (`n // k`) apart, so
    leaf `p` is `(codeword[p], codeword[p + n/k], ..., codeword[p +
    (k-1)n/k])` — the k-ary `pair_leaves`."""
    return self._regroup(codeword, self.fold_factor)

group_indices

group_indices(
    positions: Array, level: int
) -> tuple[Array, ...]

Natural order: the k-th-root coset of layer level whose fold lands at positions sits a sub-layer apart, and positions itself is the landing (first) index.

Source code in zorch/coding/reed_solomon.py

def group_indices(self, positions: Array, level: int) -> tuple[Array, ...]:
    """Natural order: the k-th-root coset of layer `level` whose fold lands
    at `positions` sits a sub-layer apart, and `positions` itself is the
    landing (first) index."""
    sub = self.block_len >> (self._log2_fold_factor * (level + 1))
    return tuple(positions + m * sub for m in range(self.fold_factor))

fold_group_values

fold_group_values(
    group: Array, beta: Array, positions: Array, level: int
) -> Array

Fold opened k-groups of layer level; the x-coordinates are the group's points on the layer's (level-times-k-th-powered) domain.

Source code in zorch/coding/reed_solomon.py

def fold_group_values(
    self, group: Array, beta: Array, positions: Array, level: int
) -> Array:
    """Fold opened k-groups of layer `level`; the x-coordinates are the
    group's points on the layer's `(level`-times-`k`-th-powered) domain."""
    n = self.block_len >> (self._log2_fold_factor * level)
    domain = eval_domain(
        _base_dtype(self.dtype), n, shift=self._group_level_shift(level)
    )
    points = domain[fnp.stack(self.group_indices(positions, level), axis=-1)]
    return self._fold_groups(group, points, beta)

group_layer_positions

group_layer_positions(
    positions: Array, num_rounds: int
) -> list[Array]

Natural order: a_i = q_i mod (n / k^{i+1}) with q_0 = positions, q_{i+1} = a_i, elementwise over the query axis — the k-ary layer_positions.

Source code in zorch/coding/reed_solomon.py

def group_layer_positions(self, positions: Array, num_rounds: int) -> list[Array]:
    """Natural order: `a_i = q_i mod (n / k^{i+1})` with `q_0 = positions`,
    `q_{i+1} = a_i`, elementwise over the query axis — the k-ary
    `layer_positions`."""
    indices = []
    q = positions
    for i in range(num_rounds):
        a = q % (self.block_len >> (self._log2_fold_factor * (i + 1)))
        indices.append(a)
        q = a
    return indices

BitReversedReedSolomon

Reed-Solomon with codewords in bit-reversed evaluation order.

Some commitment layouts store the codeword bit-reversed so a fold's point pair sits adjacently ((2p, 2p+1)) instead of a half-layer apart — Merkle paths of a pair then share all but their last node, and the layout is fold-stable (folding a bit-reversed layer yields the squared domain's codeword, again bit-reversed). The fold math is ReedSolomon's; only the layout-dependent surfaces differ — pair geometry (pair_indices / layer_positions), the fold's x-coordinate gather, and the encode/domain output order.

Source code in zorch/coding/reed_solomon.py

class BitReversedReedSolomon:
    """Reed-Solomon with codewords in bit-reversed evaluation order.

    Some commitment layouts store the codeword bit-reversed so a fold's point
    pair sits adjacently (`(2p, 2p+1)`) instead of a half-layer apart — Merkle
    paths of a pair then share all but their last node, and the layout is
    fold-stable (folding a bit-reversed layer yields the squared domain's
    codeword, again bit-reversed). The fold math is `ReedSolomon`'s; only the
    layout-dependent surfaces differ — pair geometry (`pair_indices` /
    `layer_positions`), the fold's x-coordinate gather, and the
    `encode`/`domain` output order.
    """

    def __init__(
        self,
        message_len: int,
        blowup: int,
        dtype: Any,
        *,
        coset_shift: Array | None = None,
    ) -> None:
        self._natural = ReedSolomon(message_len, blowup, dtype, coset_shift=coset_shift)
        self.message_len = message_len
        self.block_len = self._natural.block_len
        self.dtype = dtype

    # Value equality/hash via the wrapped natural-order code (see ReedSolomon);
    # the isinstance gate keeps the two layouts distinct.
    def __eq__(self, other: object) -> bool:
        if not isinstance(other, BitReversedReedSolomon):
            return NotImplemented
        return self._natural == other._natural

    def __hash__(self) -> int:
        return hash((BitReversedReedSolomon, self._natural))

    def encode(self, message: Array) -> Array:
        cw = self._natural.encode(message)
        return lax.bit_reverse(cw, dimensions=(cw.ndim - 1,))

    def domain(self) -> Array:
        """The points `encode` evaluates on, in codeword (bit-reversed) order."""
        return lax.bit_reverse(self._natural.domain(), dimensions=(0,))

    def fold(self, codeword: Array, beta: Array) -> Array:
        n = codeword.shape[0]
        if n < 2:
            raise ValueError(f"fold requires a codeword of length >= 2, got {n}")
        level = log2_strict_usize(self.block_len // n)
        pairs = codeword.reshape(n // 2, 2)
        return fri_fold_values(
            pairs[:, 0], pairs[:, 1], beta, self._pair_points(n, level)
        )

    def fold_values(
        self, lo: Array, hi: Array, beta: Array, positions: Array, level: int
    ) -> Array:
        # Gather the few queried x-coordinates straight from the natural-order
        # half at bit-reversed indices — a full-width `lax.bit_reverse` of the
        # domain would be discarded except at `positions`.
        n = self.block_len >> level
        x = self._layer_domain(n, level)[_bit_reverse_indices(positions, n // 2)]
        return fri_fold_values(lo, hi, beta, x)

    def check_final(self, final: Array, claim: Array) -> Array:
        """Constant-polynomial membership is order-invariant."""
        return self._natural.check_final(final, claim)

    def pair_indices(self, positions: Array, level: int) -> tuple[Array, Array]:
        """Bit-reversed order: the pair landing at `positions` is adjacent."""
        return positions * 2, positions * 2 + 1

    def pair_leaves(self, codeword: Array) -> Array:
        """Bit-reversed order: conjugates are adjacent, so leaf `p` is the pair
        `(codeword[2p], codeword[2p + 1])`."""
        return codeword.reshape(codeword.shape[0] // 2, 2)

    def layer_positions(self, positions: Array, num_rounds: int) -> list[Array]:
        """Bit-reversed order: each fold halves the index, `a_i = q >> (i+1)`."""
        indices = []
        q = positions
        for _ in range(num_rounds):
            q = q >> 1
            indices.append(q)
        return indices

    def _layer_domain(self, n: int, level: int) -> Array:
        """Layer `level`'s natural-order evaluation domain (length `n`)."""
        return eval_domain(
            _base_dtype(self.dtype), n, shift=self._natural._level_shift(level)
        )

    def _pair_points(self, n: int, level: int) -> Array:
        """x-coordinates of all of layer `level`'s pairs in pair order: entry
        `p` is the evaluation point of the pair `(2p, 2p+1)`, i.e. the natural
        domain's first half gathered through the bit-reversal. For a sparse
        gather see `fold_values`, which reverses the indices instead."""
        x = self._layer_domain(n, level)[: n // 2]
        if n > 2:
            x = lax.bit_reverse(x, dimensions=(0,))
        return x

domain

domain() -> Array

The points encode evaluates on, in codeword (bit-reversed) order.

Source code in zorch/coding/reed_solomon.py

def domain(self) -> Array:
    """The points `encode` evaluates on, in codeword (bit-reversed) order."""
    return lax.bit_reverse(self._natural.domain(), dimensions=(0,))

check_final

check_final(final: Array, claim: Array) -> Array

Constant-polynomial membership is order-invariant.

Source code in zorch/coding/reed_solomon.py

def check_final(self, final: Array, claim: Array) -> Array:
    """Constant-polynomial membership is order-invariant."""
    return self._natural.check_final(final, claim)

pair_indices

pair_indices(
    positions: Array, level: int
) -> tuple[Array, Array]

Bit-reversed order: the pair landing at positions is adjacent.

Source code in zorch/coding/reed_solomon.py

def pair_indices(self, positions: Array, level: int) -> tuple[Array, Array]:
    """Bit-reversed order: the pair landing at `positions` is adjacent."""
    return positions * 2, positions * 2 + 1

pair_leaves

pair_leaves(codeword: Array) -> Array

Bit-reversed order: conjugates are adjacent, so leaf p is the pair (codeword[2p], codeword[2p + 1]).

Source code in zorch/coding/reed_solomon.py

def pair_leaves(self, codeword: Array) -> Array:
    """Bit-reversed order: conjugates are adjacent, so leaf `p` is the pair
    `(codeword[2p], codeword[2p + 1])`."""
    return codeword.reshape(codeword.shape[0] // 2, 2)

layer_positions

layer_positions(
    positions: Array, num_rounds: int
) -> list[Array]

Bit-reversed order: each fold halves the index, a_i = q >> (i+1).

Source code in zorch/coding/reed_solomon.py

def layer_positions(self, positions: Array, num_rounds: int) -> list[Array]:
    """Bit-reversed order: each fold halves the index, `a_i = q >> (i+1)`."""
    indices = []
    q = positions
    for _ in range(num_rounds):
        q = q >> 1
        indices.append(q)
    return indices

eval_domain

eval_domain(
    dtype: Any,
    n: int,
    *,
    shift: Array | None = None,
    generator: int | None = None
) -> Array

The order-n two-adic subgroup points [d₀..d_{n-1}] in lax.ntt order, or the coset points [shift·d₀..shift·d_{n-1}] when shift is given.

lax.ntt of the coefficient vector of p(X)=X (i.e. e₁) returns [p(d₀)..p(d_{n-1})] = [d₀..d_{n-1}], so the domain is read off the same NTT the encoder uses. generator selects the subgroup generator the NTT root is gen^((p-1)/n) of; None uses the dtype's canonical root. It must match the generator the codeword was encoded with — the domain order is the root's. n must be a power of two; the order-1 subgroup is {1}.

Source code in zorch/coding/reed_solomon.py

def eval_domain(
    dtype: Any, n: int, *, shift: Array | None = None, generator: int | None = None
) -> Array:
    """The order-`n` two-adic subgroup points [d₀..d_{n-1}] in `lax.ntt` order,
    or the coset points [shift·d₀..shift·d_{n-1}] when `shift` is given.

    `lax.ntt` of the coefficient vector of p(X)=X (i.e. e₁) returns
    [p(d₀)..p(d_{n-1})] = [d₀..d_{n-1}], so the domain is read off the same NTT
    the encoder uses. `generator` selects the subgroup generator the NTT root is
    `gen^((p-1)/n)` of; None uses the dtype's canonical root. It must match the
    `generator` the codeword was encoded with — the domain order is the root's.
    `n` must be a power of two; the order-1 subgroup is {1}."""
    if not is_power_of_two(n):
        raise ValueError(f"eval_domain size must be a power of two, got {n}")
    if n == 1:
        domain = fnp.ones((1,), dtype)
    else:
        e1 = fnp.zeros(n, dtype).at[1].set(fnp.ones((), dtype))
        domain = lax.ntt(e1, ntt_type="NTT", ntt_length=n, generator=generator)
    return domain if shift is None else shift * domain

fri_fold_values

fri_fold_values(
    fx: Array, fnx: Array, beta: Array, x: Array
) -> Array

g(x²) = (f(x)+f(−x))/2 + β·(f(x)−f(−x))/(2x). f-values may be EF; x carries the domain's dtype.

Source code in zorch/coding/reed_solomon.py

def fri_fold_values(fx: Array, fnx: Array, beta: Array, x: Array) -> Array:
    """g(x²) = (f(x)+f(−x))/2 + β·(f(x)−f(−x))/(2x). f-values may be EF; x carries
    the domain's dtype."""
    one = fnp.ones((), fx.dtype)
    two = one + one
    return (fx + fnx) / two + beta * (fx - fnx) / (two * x)

fri_fold_k

fri_fold_k(
    group: Array,
    beta: Array,
    *,
    points: Array | None = None,
    coset: tuple[Array, Array] | None = None
) -> Array

k-ary FRI fold: the degree-(k-1) interpolant through a group's k points, evaluated at beta. Pass exactly one of:

• points (..., k): arbitrary coordinates in the caller's own domain order — per-group Lagrange, no domain convention. The per-query verifier path.
• coset = (coset_inv, generator): the points form a coset s·⟨ω⟩ — batched lax.ntt INTT, evaluated at coset_inv·beta (the unshift folds into the point). coset_inv is per-group s⁻¹; generator selects ω as generator^((p-1)/k), None for the canonical root. The prover path.

Byte-identical wherever both apply (the interpolant is unique); values and coordinates may be extension-field.

Source code in zorch/coding/reed_solomon.py

def fri_fold_k(
    group: Array,
    beta: Array,
    *,
    points: Array | None = None,
    coset: tuple[Array, Array] | None = None,
) -> Array:
    """k-ary FRI fold: the degree-`(k-1)` interpolant through a group's `k`
    points, evaluated at `beta`. Pass exactly one of:

    - `points` (..., k): arbitrary coordinates in the caller's own domain order
      — per-group Lagrange, no domain convention. The per-query verifier path.
    - `coset = (coset_inv, generator)`: the points form a coset `s·⟨ω⟩` —
      batched `lax.ntt` INTT, evaluated at `coset_inv·beta` (the unshift folds
      into the point). `coset_inv` is per-group `s⁻¹`; `generator` selects ω as
      `generator^((p-1)/k)`, None for the canonical root. The prover path.

    Byte-identical wherever both apply (the interpolant is unique); values and
    coordinates may be extension-field."""
    if (points is None) == (coset is None):
        raise ValueError("fri_fold_k needs exactly one of `points` or `coset`")
    if coset is not None:
        coset_inv, generator = coset
        coeffs = lax.ntt(
            group, ntt_type="INTT", ntt_length=group.shape[-1], generator=generator
        )
        return eval_coeffs(coeffs, coset_inv * beta)
    # Lagrange path: unroll the linear combination over the static factor k
    # rather than `fnp.dot` — a reduction is a kInput/gather fusion boundary on
    # GPU, so the fold would not lower to one fused kernel (CLAUDE.md "Fusion by
    # construction"). k is a small compile-time constant, so the unroll is cheap.
    basis = compute_lagrange_basis(beta, points)
    folded = group[..., 0] * basis[..., 0]
    for i in range(1, group.shape[-1]):
        folded = folded + group[..., i] * basis[..., i]
    return folded

fri_fold

fri_fold(
    codeword: Array,
    beta: Array,
    *,
    shift: Array | None = None
) -> Array

FRI-fold a natural-order RS codeword (length 2^m) by β, halving its length.

Natural order: dⱼ and d_{j+n/2} = −dⱼ are conjugates, so f(x)=codeword[:half], f(−x)=codeword[half:], x=domain[:half]. Result is the fold over the order-(n/2) squared domain, again in natural order.

shift is the coset shift of the codeword's own domain. The fold lands on the squared domain, so the next layer's shift is shift² — iterating callers must square it each round.

Source code in zorch/coding/reed_solomon.py

def fri_fold(codeword: Array, beta: Array, *, shift: Array | None = None) -> Array:
    """FRI-fold a natural-order RS codeword (length 2^m) by β, halving its length.

    Natural order: dⱼ and d_{j+n/2} = −dⱼ are conjugates, so f(x)=codeword[:half],
    f(−x)=codeword[half:], x=domain[:half]. Result is the fold over the order-(n/2)
    squared domain, again in natural order.

    `shift` is the coset shift of the codeword's own domain. The fold lands on
    the squared domain, so the next layer's shift is `shift²` — iterating
    callers must square it each round."""
    n = codeword.shape[0]
    if n < 2:
        raise ValueError(f"fri_fold requires a codeword of length >= 2, got {n}")
    half = n // 2
    domain = eval_domain(_base_dtype(codeword.dtype), n, shift=shift)
    return fri_fold_values(codeword[:half], codeword[half:], beta, domain[:half])