zorch.pcs.jagged.branching_program

Jagged indicator via a 4-state branching program.

The MLE of the jagged indicator H(r, c, i) = 1 iff i = t_c + r and t_c <= i < t_{c+1} evaluated at the sumcheck's random point, via a 4-state (carry, comparison) automaton folded over the prefix bits MSB->LSB — the "un-jagging" glue that reduces the ragged trace to a dense sumcheck.

bp_eval_core is the per-column fold, wrapped in the zorch.jagged_bp name-routed composite so a vendor fuses the whole DP into one register-resident kernel. The branching program exists only for this indicator (no other caller), so the marker is domain-named rather than a generic matrix_fold.

bp_eval_core

bp_eval_core(
    z_row: Array,
    z_index: Array,
    prefix_sum: Array,
    next_prefix_sum: Array,
    t_matrix: Array,
) -> Array

h(z_row, z_index; t_c, t_{c+1}) — the per-column jagged BP indicator eval, wrapped in the zorch.jagged_bp name-routed composite so a vendor fuses the whole DP fold into one register-resident kernel (the CLAUDE.md fusion non-negotiable). The decomposition is byte-identical, so an unrecognized marker lowers with no behavior change. A vmap over columns (_bp_all) batches this single-column composite, mirroring Poseidon2's single-state permute marker.

The layer count is NOT an operand: it is derived from the operand shapes (max(z_row.shape[0], prefix_sum.shape[0])) in the decomposition, since a bare scalar operand is constant-sunk into the fusion and would break the emitter's positional ABI. Deriving keeps a symbolic export dim symbolic.

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

def bp_eval_core(
    z_row: Array,
    z_index: Array,
    prefix_sum: Array,
    next_prefix_sum: Array,
    t_matrix: Array,
) -> Array:
    """h(z_row, z_index; t_c, t_{c+1}) — the per-column jagged BP indicator eval,
    wrapped in the `zorch.jagged_bp` name-routed composite so a vendor fuses the
    whole DP fold into one register-resident kernel (the CLAUDE.md fusion
    non-negotiable). The decomposition is byte-identical, so an unrecognized marker
    lowers with no behavior change. A ``vmap`` over columns (`_bp_all`) batches
    this single-column composite, mirroring Poseidon2's single-state permute
    marker.

    The layer count is NOT an operand: it is derived from the operand shapes
    (``max(z_row.shape[0], prefix_sum.shape[0])``) in the decomposition, since a
    bare scalar operand is constant-sunk into the fusion and would break the
    emitter's positional ABI. Deriving keeps a symbolic export dim symbolic.
    """
    return fused_region(
        _bp_eval_decomposition,
        z_row,
        z_index,
        prefix_sum,
        next_prefix_sum,
        t_matrix,
        name=JAGGED_BP_MARKER,
        version=JAGGED_BP_MARKER_VERSION,
        num_memory_states=NUM_MEMORY_STATES,
        num_bit_states=NUM_BIT_STATES,
    )