zorch.sumcheck.eq.accumulators

Accumulator pre-computation for the small-value EqPoly sumcheck (Procedure 9).

Each round i ≤ l₀ of the small-value phase reads a precomputed table Aᵢ(v, u) that already contracts every variable except the first i, so the round polynomial is a cheap contraction Rᵢ·Aᵢ instead of a fresh pass over the cube. Aᵢ(v, u) sums Πₖ pₖ(v, u, x) weighted by eq(w, ·) over the extended prefix v ∈ U_dⁱ⁻¹, node u ∈ Û_d, and the bound suffix. precompute_accumulators fills all l₀ tables in one sweep over β ∈ U_dˡ⁰; the index maps below route each β to the (round, v, u, y) slots it contributes to.

precompute_accumulators

precompute_accumulators(
    p_evals: Array, e_in: Array, e_out: list[Array]
) -> list[Array]

Tables [A₁, …, Aₗ₀], Aᵢ of shape ((d+1)ⁱ⁻¹, d) over (v ∈ U_dⁱ⁻¹, u ∈ Û_d). e_in = eq(w_in, x_in) over {0,1}^{l/2}; e_out[i−1] = eq over (y, x_out) for round i (both the y and x_out domains, not x_out alone).

Source code in zorch/sumcheck/eq/accumulators.py

def precompute_accumulators(
    p_evals: Array, e_in: Array, e_out: list[Array]
) -> list[Array]:
    """Tables [A₁, …, Aₗ₀], Aᵢ of shape ((d+1)ⁱ⁻¹, d) over (v ∈ U_dⁱ⁻¹, u ∈ Û_d).
    e_in = eq(w_in, x_in) over {0,1}^{l/2}; e_out[i−1] = eq over (y, x_out) for
    round i (both the y and x_out domains, not x_out alone)."""
    d = p_evals.shape[0]
    l = log2_strict_usize(p_evals.shape[1])
    l_half = l // 2
    l_0 = len(e_out)
    x_in_size = 1 << l_half
    x_out_size = 1 << (l - l_half - l_0)
    beta_size = (d + 1) ** l_0

    # Extend the first l₀ variables to U_d, then contract x_in against eq(w_in, ·):
    # tA[β, x_out] = Σ_{x_in} e_in[x_in] · Πₖ pₖ(β, x_in, x_out).
    p_beta = vmap(lambda e: _extend_prefix_to_domain(e, d, l_0))(p_evals)
    p_beta_prod = fnp.reshape(
        fnp.prod(p_beta, axis=0), (beta_size, x_in_size, x_out_size)
    )
    t_a = (p_beta_prod * e_in[None, :, None]).sum(axis=1)  # (β, x_out)

    # _idx4 is pure host arithmetic, so collect each round's (β, flat v·d+u, y) slots
    # once on the host, then contract + scatter that round in a single vectorized op.
    # A per-contribution `.at[].add()` instead bakes O((d+1)^l₀·l₀) scatters into the
    # traced graph — the compile time/memory blow up exponentially in l₀.
    routes: list[tuple[list[int], list[int], list[int]]] = [
        ([], [], []) for _ in range(l_0)
    ]
    for beta_idx in range(beta_size):
        for i, v, u, y in _idx4(beta_idx, l_0, d):
            betas, vus, ys = routes[i - 1]
            betas.append(beta_idx)
            vus.append(v * d + u)
            ys.append(y)

    accumulators = []
    for i in range(l_0):
        betas, vus, ys = (fnp.asarray(col, dtype=fnp.int32) for col in routes[i])
        # contrib[k] = Σ_x_out e_out[i][y_k, :] · t_a[β_k, :]
        contrib = fnp.sum(e_out[i][ys, :] * t_a[betas, :], axis=1)
        flat = fnp.zeros((d + 1) ** i * d, dtype=p_evals.dtype).at[vus].add(contrib)
        accumulators.append(flat.reshape(((d + 1) ** i, d)))
    return accumulators