zorch.pcs.jagged.poly

Jagged Little Polynomial — layout construction + polymorphic partial-eval.

Builds the column prefix-sum bit tensor (build_jagged_layout) and evaluates the partial MLE J̃(z_row, z_col, ·) over the dense area (partial_eval_core), shape-polymorphic in the column count and prefix-bit width, in AOT-clean form: a static l_max column axis, n_d bound to the instance's log-area tier. The branching-program indicator eval that closes the sumcheck lives in branching_program.py.

msb_first_bits

msb_first_bits(values: Any, num_bits: int) -> np.ndarray

(N,) ints → (N, num_bits) numpy int64, MSB first. Host-side; never feeds a field element into >> (XLA field dtypes have no lax.shift).

Source code in zorch/pcs/jagged/poly.py

def msb_first_bits(values: Any, num_bits: int) -> np.ndarray:
    """(N,) ints → (N, num_bits) numpy int64, MSB first. Host-side; never feeds a
    field element into >> (XLA field dtypes have no lax.shift)."""
    arr = np.asarray(values, dtype=np.int64)
    shifts = np.arange(num_bits - 1, -1, -1, dtype=np.int64)
    return (arr[:, None] >> shifts[None, :]) & 1

build_jagged_layout

build_jagged_layout(
    row_counts: Sequence[int], l_max: int, dtype: Any
) -> tuple[Array, int]

heights -> (col_prefix_sums (l_max+1, n_d) field bit tensor, n_d).

n_d (the BP layer count = log-area tier) is the only static dim downstream code needs; every other dim is derived from array shapes by the shape-polymorphic cores. Padding columns use an EMPTY RANGE (t_c = t_{c+1} = t_L) — zero-bit padding injects a phantom range that corrupts J̃, so it is forbidden.

Source code in zorch/pcs/jagged/poly.py

def build_jagged_layout(
    row_counts: Sequence[int], l_max: int, dtype: Any
) -> tuple[Array, int]:
    """heights -> (col_prefix_sums (l_max+1, n_d) field bit tensor, n_d).

    ``n_d`` (the BP layer count = log-area tier) is the only static dim downstream
    code needs; every other dim is derived from array shapes by the shape-polymorphic
    cores. Padding columns use an EMPTY RANGE (t_c = t_{c+1} = t_L) — zero-bit
    padding injects a phantom range that corrupts J̃, so it is forbidden.
    """
    real_L = len(row_counts)
    if real_L > l_max:
        raise ValueError(f"real_L={real_L} > l_max={l_max}")
    prefix = build_prefix_sums(row_counts)  # length real_L+1
    n_d = log_area_tier(prefix[-1])
    padded = prefix + [prefix[-1]] * (l_max - real_L)  # length l_max+1, empty-range pad
    cps = fnp.asarray(msb_first_bits(padded, n_d), dtype=dtype)
    return cps, n_d

partial_eval_core

partial_eval_core(
    col_prefix_sums: Array,
    z_row: Array,
    z_col: Array,
    size: Any,
) -> Array

J tilde over the first size dense indices -- the jagged outer indicator the outer Hadamard sumcheck folds against D.

Decodes the prefix tensor and builds the column-eq table (both cheap, kept outside the marker), then wraps the 2ⁿᵈ scatter in the jagged_indicator composite for a vendor to fuse. Heights ride as data, so the compile keys only on the n_d class, never the per-column heights.

Source code in zorch/pcs/jagged/poly.py

def partial_eval_core(
    col_prefix_sums: Array,
    z_row: Array,
    z_col: Array,
    size: Any,
) -> Array:
    """J tilde over the first size dense indices -- the jagged outer indicator the
    outer Hadamard sumcheck folds against D.

    Decodes the prefix tensor and builds the column-eq table (both cheap, kept
    outside the marker), then wraps the 2ⁿᵈ scatter in the jagged_indicator
    composite for a vendor to fuse. Heights ride as data, so the compile keys only
    on the n_d class, never the per-column heights."""
    n_d = col_prefix_sums.shape[1]
    prefix_sums_int = _decode_prefix_sums(col_prefix_sums, n_d)
    col_eq = expand_eq_to_hypercube(z_col, fnp.ones([], dtype=z_row.dtype))
    # Search depth from the real column count when static, so a cap far above 2ⁿᶜ
    # is searched fully; +2 keeps l_max==1 safe. Symbolic L has no static length, so
    # it falls back to the n_c+2 proxy.
    ncols = col_prefix_sums.shape[0]
    search_steps = (
        log2_ceil_usize(ncols) + 2 if isinstance(ncols, int) else z_col.shape[0] + 2
    )
    return fused_region(
        _partial_eval_decomposition,
        prefix_sums_int,
        col_eq,
        z_row,
        name=JAGGED_INDICATOR_MARKER,
        version=JAGGED_INDICATOR_MARKER_VERSION,
        size=size,
        search_steps=search_steps,
    )