src/hamiltonzero/evaluation/sequence_trunk.py

# Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0


from __future__ import annotations

from collections.abc import Callable
from functools import partial
from math import gcd

import equinox as eqx
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

from hamiltonzero.compiled.types import SharedTrunk, TrunkCompilerKernel
from hamiltonzero.model.fused_silu import fused_silu
from hamiltonzero.model.global_ladder import (
    BoundaryGlobalUpdate,
    ResidualGlobalUpdate,
    TreeGlobalUpdate,
)
from hamiltonzero.model.readout_leaf_context import (
    PhysicalReadoutContext,
    PhysicalReadoutContextLayer,
    RouterContext,
    RouterContextLayer,
    lca_alibi_bias,
    lca_fixed_slopes,
)
from .sequence_parallel import (
    pallas_rectangular_edge_attention,
    ring_learned_fwl2_columns_local,
    ring_learned_fwl2_local,
)
from hamiltonzero.model.tree import _tree_ngpt_residual, _tree_sphere


def _local_rows(x: jax.Array, *, axis_name: str, local_size: int) -> jax.Array:
    start = jax.lax.axis_index(axis_name) * local_size
    return jax.lax.dynamic_slice_in_dim(x, start, local_size, axis=0)


def _global_row_indices(*, axis_name: str, local_size: int) -> jax.Array:
    start = jax.lax.axis_index(axis_name) * local_size
    return start + jnp.arange(local_size, dtype=jnp.int32)


def _gather_rows(x: jax.Array, *, axis_name: str) -> jax.Array:

    return jax.lax.all_gather(x, axis_name=axis_name, axis=0, tiled=True)


def ring_permute_rows_local(
    values: jax.Array,
    permutation: jax.Array,
    *,
    axis_name: str,
    axis_size: int,
) -> jax.Array:

    local_size = values.shape[0]
    lane = jax.lax.axis_index(axis_name).astype(jnp.int32)
    output_ids = jax.lax.dynamic_slice_in_dim(
        permutation,
        lane * local_size,
        local_size,
        axis=0,
    ).astype(jnp.int32)
    owners = output_ids // local_size
    offsets = output_ids % local_size
    output = jnp.zeros_like(values)

    def select(panel, origin, current):
        selected = panel[offsets]
        take = owners == origin
        while take.ndim < selected.ndim:
            take = take[..., None]
        return jnp.where(take, selected, current)

    origin = lane
    output = select(values, origin, output)
    ring = tuple((i, (i + 1) % axis_size) for i in range(axis_size))

    def step(carry, _):
        panel, panel_origin, current = carry
        panel = jax.lax.ppermute(panel, axis_name, ring)
        panel_origin = (panel_origin - jnp.asarray(1, dtype=jnp.int32)) % axis_size
        return (panel, panel_origin, select(panel, panel_origin, current)), None

    (_, _, output), _ = jax.lax.scan(
        step,
        (values, origin, output),
        xs=None,
        length=axis_size - 1,
    )
    return output


def permute_pair_rows_and_columns_local(
    edge_rows: jax.Array,
    permutation: jax.Array,
    *,
    axis_name: str,
    axis_size: int,
) -> jax.Array:

    rows = ring_permute_rows_local(
        edge_rows,
        permutation,
        axis_name=axis_name,
        axis_size=axis_size,
    )
    return rows[:, permutation]


def transpose_pair_rows_local(
    edge_rows: jax.Array,
    *,
    axis_name: str,
    axis_size: int,
) -> jax.Array:

    if axis_size == 1:
        return jnp.swapaxes(edge_rows, 0, 1)
    local_size, n = edge_rows.shape[:2]
    if n != local_size * axis_size:
        raise ValueError(
            "pair transpose requires N == local_rows * axis_size, got "
            f"N={n}, local_rows={local_size}, axis_size={axis_size}"
        )
    trailing = edge_rows.shape[2:]
    send = edge_rows.reshape((local_size, axis_size, local_size) + trailing)
    send = jnp.swapaxes(send, 0, 1)
    received = jax.lax.all_to_all(
        send,
        axis_name,
        split_axis=0,
        concat_axis=0,
    )
    received = jnp.transpose(
        received,
        (2, 0, 1) + tuple(range(3, received.ndim)),
    )
    return received.reshape((local_size, n) + trailing)


def _linear(module, x: jax.Array) -> jax.Array:
    out = x @ module.weight
    bias = getattr(module, "bias", None)
    return out if bias is None else out + bias


def _norm(module, x: jax.Array) -> jax.Array:

    weight = getattr(module, "weight", None)
    if weight is None:
        return x
    stats_dtype = jnp.promote_types(jnp.float32, x.dtype)
    x_hi = x.astype(stats_dtype) if x.dtype != stats_dtype else x
    bias = getattr(module, "bias", None)
    if bias is None:
        stat = jnp.mean(x_hi * x_hi, axis=-1, keepdims=True)
        normalized = x_hi * jax.lax.rsqrt(stat + module.eps)
    else:
        mean = jnp.mean(x_hi, axis=-1, keepdims=True)
        centered = x_hi - mean
        stat = jnp.mean(centered * centered, axis=-1, keepdims=True)
        normalized = centered * jax.lax.rsqrt(stat + module.eps)
    normalized = normalized.astype(x.dtype)
    out = normalized * weight.astype(x.dtype)
    return out if bias is None else out + bias.astype(x.dtype)


def _raw_norm(x, scale, shift, *, eps: float):
    stats_dtype = jnp.promote_types(jnp.float32, x.dtype)
    x_hi = x.astype(stats_dtype)
    if shift is None:
        normalized = x_hi * jax.lax.rsqrt(
            jnp.mean(x_hi * x_hi, axis=-1, keepdims=True) + eps
        )
    else:
        centered = x_hi - jnp.mean(x_hi, axis=-1, keepdims=True)
        normalized = centered * jax.lax.rsqrt(
            jnp.mean(centered * centered, axis=-1, keepdims=True) + eps
        )
    out = normalized.astype(x.dtype) * scale.astype(x.dtype)
    return out if shift is None else out + shift.astype(x.dtype)


def _mlp_after_input_projection(mlp, hidden: jax.Array) -> jax.Array:
    for norm, l1, l2 in zip(mlp.block_norms, mlp.block_l1s, mlp.block_l2s, strict=True):
        hidden = hidden + mlp.inner_gain * _linear(
            l2, mlp._act(_linear(l1, _norm(norm, hidden)))
        )
    return _linear(mlp.out_proj, _norm(mlp.out_norm, hidden))


def _mlp(mlp, x: jax.Array) -> jax.Array:
    return _mlp_after_input_projection(mlp, _linear(mlp.in_proj, x))


def _unnormalized_mlp(mlp, x: jax.Array) -> jax.Array:
    hidden = _linear(mlp.in_proj, x)
    for l1, l2 in zip(mlp.block_l1s, mlp.block_l2s, strict=True):
        hidden = hidden + mlp.inner_gain * _linear(l2, mlp._act(_linear(l1, hidden)))
    return _linear(mlp.out_proj, hidden)


def _split_linear(module, parts: tuple[jax.Array, ...]) -> jax.Array:

    offset = 0
    out = module.bias
    for part in parts:
        width = part.shape[-1]
        out = out + part @ module.weight[offset : offset + width]
        offset += width
    if offset != module.weight.shape[0]:
        raise ValueError(
            f"split input width {offset} does not match weight {module.weight.shape[0]}"
        )
    return out


def _mlp_split_input(mlp, parts: tuple[jax.Array, ...]) -> jax.Array:
    return _mlp_after_input_projection(mlp, _split_linear(mlp.in_proj, parts))


def _g_descriptor_pool(pool, g, xs, mask):

    n = xs.shape[0]
    xn = _norm(pool.ln_in, xs)
    q = (g @ pool.W_q).reshape(pool.n_heads, pool.d_k)
    k = _linear(pool.K, xn).reshape(n, pool.n_heads, pool.d_k)
    v = _linear(pool.V, xn).reshape(n, pool.n_heads, pool.d_v)
    scores = jnp.einsum("hd,nhd->hn", q, k) / jnp.sqrt(
        jnp.asarray(pool.d_k, dtype=xs.dtype)
    )
    scores = jnp.where(
        mask[None, :] > 0,
        scores,
        jnp.asarray(-1.0e30, dtype=scores.dtype),
    )
    weights = jax.nn.softmax(scores, axis=-1)
    return jnp.einsum("hn,nhv->hv", weights, v).reshape(-1)


def _g_descriptor_pool_rows(pool, g, xs_rows, key_mask, *, tile_size: int = 128):

    r, n = xs_rows.shape[:2]
    if tile_size < 1:
        raise ValueError("descriptor-pool tile_size must be positive")
    tile_width = gcd(n, min(n, int(tile_size)))
    tile_count = n // tile_width
    q = (g @ pool.W_q).reshape(pool.n_heads, pool.d_k)
    scale = jnp.sqrt(jnp.asarray(pool.d_k, dtype=xs_rows.dtype))
    scores0 = jnp.zeros((r, pool.n_heads, n), dtype=xs_rows.dtype)

    def score_tile(tile_index, scores):
        start = tile_index * tile_width
        xs_tile = jax.lax.dynamic_slice_in_dim(xs_rows, start, tile_width, axis=1)
        xn_tile = _norm(pool.ln_in, xs_tile)
        k_tile = _linear(pool.K, xn_tile).reshape(r, tile_width, pool.n_heads, pool.d_k)
        tile_scores = jnp.einsum("hd,rthd->rht", q, k_tile) / scale
        mask_tile = jax.lax.dynamic_slice_in_dim(key_mask, start, tile_width, axis=0)
        tile_scores = jnp.where(
            mask_tile[None, None, :] > 0,
            tile_scores,
            jnp.asarray(-1.0e30, dtype=tile_scores.dtype),
        )
        return jax.lax.dynamic_update_slice_in_dim(scores, tile_scores, start, axis=2)

    scores = jax.lax.fori_loop(0, tile_count, score_tile, scores0)
    weights = jax.nn.softmax(scores, axis=-1)
    pooled0 = jnp.zeros((r, pool.n_heads, pool.d_v), dtype=xs_rows.dtype)

    def value_tile(tile_index, pooled):
        start = tile_index * tile_width
        xs_tile = jax.lax.dynamic_slice_in_dim(xs_rows, start, tile_width, axis=1)
        xn_tile = _norm(pool.ln_in, xs_tile)
        v_tile = _linear(pool.V, xn_tile).reshape(r, tile_width, pool.n_heads, pool.d_v)
        weight_tile = jax.lax.dynamic_slice_in_dim(weights, start, tile_width, axis=2)
        return pooled + jnp.einsum("rht,rthv->rhv", weight_tile, v_tile)

    pooled = jax.lax.fori_loop(0, tile_count, value_tile, pooled0)
    return pooled.reshape(r, -1)


def _g_update(update, g, pooled):
    g_input = g @ update.g_tap_w
    x = jnp.concatenate((g_input, pooled.astype(g.dtype)))
    stats = jnp.mean(jnp.square(x), keepdims=True)
    x = x * jax.lax.rsqrt(stats + 1.0e-5) * update.ln_s
    hidden = fused_silu(x @ update.w1 + update.b1)
    delta = hidden @ update.w2 + update.b2
    if isinstance(update, ResidualGlobalUpdate):
        return g + update.residual_gain * delta
    if isinstance(update, BoundaryGlobalUpdate):
        return _tree_sphere(g + delta)
    if not isinstance(update, TreeGlobalUpdate):
        raise TypeError(f"unsupported global update {type(update)!r}")
    skip = _tree_sphere(g)
    proposal = _tree_sphere(delta)
    gain = update.alpha_max * jax.nn.sigmoid(update.alpha)
    return _tree_sphere(skip + gain * (proposal - skip))


def _edge_row_col_global_update(
    module,
    g,
    edge_rows,
    mask,
    row_mask,
    *,
    axis_name: str,
    tile_size: int = 128,
):
    row_desc_rows = _g_descriptor_pool_rows(
        module.row_pool, g, edge_rows, mask, tile_size=tile_size
    )
    row_desc = _gather_rows(row_desc_rows, axis_name=axis_name)

    edge_column_rows = transpose_pair_rows_local(
        edge_rows,
        axis_name=axis_name,
        axis_size=edge_rows.shape[1] // edge_rows.shape[0],
    )
    col_desc_rows = _g_descriptor_pool_rows(
        module.col_pool, g, edge_column_rows, mask, tile_size=tile_size
    )
    col_desc = _gather_rows(col_desc_rows, axis_name=axis_name)
    descriptors = jnp.concatenate((row_desc, col_desc), axis=0)
    descriptor_mask = jnp.concatenate((mask, mask), axis=0)
    pooled = _g_descriptor_pool(module.set_pool, g, descriptors, descriptor_mask)
    return _g_update(module.update, g, pooled)


def sequence_parallel_edge_global_update_local(
    module,
    g,
    edge_rows,
    mask,
    *,
    axis_name: str,
    tile_size: int = 128,
):

    row_indices = _global_row_indices(
        axis_name=axis_name, local_size=edge_rows.shape[0]
    )
    return _edge_row_col_global_update(
        module,
        g,
        edge_rows,
        mask,
        mask[row_indices],
        axis_name=axis_name,
        tile_size=tile_size,
    )


def _psi_project(edge_update, parts, *, left: bool):
    if left:
        linear_in = edge_update.psi_L_in
        linear_out = edge_update.psi_L_out
    else:
        linear_in = edge_update.psi_R_in
        linear_out = edge_update.psi_R_out
    value = _split_linear(linear_in, parts)
    hidden = fused_silu(value)
    return _linear(linear_out, hidden)


def _context_edge_update(
    edge_update,
    edge_rows,
    even_rows,
    even_all,
    mask,
    row_mask,
    *,
    axis_name: str,
    axis_size: int,
):

    edge_ln = _norm(edge_update.ln_edge, edge_rows)
    even_rows_ln = _norm(edge_update.ln_even, even_rows)
    even_all_ln = _norm(edge_update.ln_even, even_all)
    if edge_update.node_ctx_proj is not None:
        even_rows_ln = _linear(edge_update.node_ctx_proj, even_rows_ln)
        even_all_ln = _linear(edge_update.node_ctx_proj, even_all_ln)
    row_endpoint = even_rows_ln[:, None, :]
    column_endpoint = even_all_ln[None, :, :]

    pair_parts = (edge_ln, row_endpoint, column_endpoint)
    left = _psi_project(edge_update, pair_parts, left=True)
    right = _psi_project(edge_update, pair_parts, left=False)
    left = left * (
        row_mask[:, None, None].astype(left.dtype)
        * mask[None, :, None].astype(left.dtype)
    )
    right = right * mask[None, :, None].astype(right.dtype)
    path = ring_learned_fwl2_local(
        left,
        right,
        axis_name=axis_name,
        axis_size=axis_size,
    )
    n_eff = jnp.maximum(jnp.sum(mask), 1.0).astype(path.dtype)
    path = _norm(edge_update.ln_path, path / jnp.sqrt(n_eff))
    return _mlp_split_input(edge_update.ffn, (*pair_parts, path))


def _rectangular_block_attention(
    attention,
    even_rows,
    edge_rows,
    mask,
    *,
    axis_name: str,
    block_k: int,
):

    r = even_rows.shape[0]
    qkv_rows = _linear(attention.W_QKV, even_rows).reshape(
        r, 3, attention.n_heads_kernel, attention.d_head
    )
    qkv_all = _gather_rows(qkv_rows, axis_name=axis_name)
    query = qkv_rows[:, 0]
    key = qkv_all[:, 1]
    value = qkv_all[:, 2]

    edge_pre = _norm(attention.ln_edge, edge_rows)
    bias = _unnormalized_mlp(attention.bias_mlp, edge_pre)
    bias = bias / jnp.sqrt(jnp.asarray(attention.d_head, bias.dtype))
    out = pallas_rectangular_edge_attention(
        query,
        key,
        value,
        bias,
        mask,
        block_k=block_k,
    )
    gate = out[:, : attention.n_heads]
    value_out = out[:, attention.n_heads :]
    out = jax.nn.sigmoid(gate) * value_out
    return _linear(attention.W_O, out.reshape(r, -1))


def _sequence_transformer_block(
    block,
    even_rows,
    edge_rows,
    g,
    mask,
    row_mask,
    *,
    axis_name: str,
    axis_size: int,
    attention_block_k: int,
):
    even_all = _gather_rows(even_rows, axis_name=axis_name)
    edge_delta = _context_edge_update(
        block.edge_update_ctx,
        edge_rows,
        even_rows,
        even_all,
        mask,
        row_mask,
        axis_name=axis_name,
        axis_size=axis_size,
    )
    edge_rows = edge_rows + block.residual_gain * edge_delta

    even_pre = _norm(block.ln_attn, even_rows)
    attention_delta = _rectangular_block_attention(
        block.attn,
        even_pre,
        edge_rows,
        mask,
        axis_name=axis_name,
        block_k=attention_block_k,
    )
    even_rows = even_rows + block.residual_gain * attention_delta

    even_pre = _norm(block.ln_ffn, even_rows)
    even_pre = even_pre + (g @ block.g_ffn_proj_w)[None].astype(even_pre.dtype)
    ffn_delta = _linear(block.ffn.l2, fused_silu(_linear(block.ffn.l1, even_pre)))
    even_rows = even_rows + block.residual_gain * ffn_delta

    even_all = _gather_rows(even_rows, axis_name=axis_name)
    pooled = _g_descriptor_pool(block.g_pool, g, even_all, mask)
    g = _g_update(block.g_update, g, pooled)
    return even_rows, edge_rows, g


def _sequence_trunk_local(
    trunk,
    local_rows,
    edge_rows,
    g,
    mask,
    row_mask,
    *,
    axis_name: str,
    axis_size: int,
    attention_block_k: int,
):
    step = partial(
        _sequence_transformer_block,
        mask=mask,
        row_mask=row_mask,
        axis_name=axis_name,
        axis_size=axis_size,
        attention_block_k=attention_block_k,
    )
    dynamic, static = eqx.partition(trunk.blocks, eqx.is_array)

    def scan_step(carry, layer_dynamic):
        block = eqx.combine(layer_dynamic, static)
        return step(block, *carry), None

    (local_rows, edge_rows, g), _ = jax.lax.scan(
        scan_step, (local_rows, edge_rows, g), dynamic
    )
    return local_rows, edge_rows, g


def _pair_mlp_tiled(mlp, values, *, tile_size: int):

    n = values.shape[1]
    tile_width = min(n, int(tile_size))
    if tile_width < 1:
        raise ValueError("pair MLP tile_size must be positive")
    full_tiles = n // tile_width
    tail_start = full_tiles * tile_width
    output = jnp.zeros(
        values.shape[:2] + (int(mlp.out_proj.weight.shape[1]),),
        dtype=values.dtype,
    )

    def update_tile(start, width, current):
        tile = jax.lax.dynamic_slice_in_dim(values, start, width, axis=1)
        return jax.lax.dynamic_update_slice_in_dim(
            current, _mlp(mlp, tile), start, axis=1
        )

    output = jax.lax.fori_loop(
        0,
        full_tiles,
        lambda tile_index, current: update_tile(
            tile_index * tile_width, tile_width, current
        ),
        output,
    )
    if tail_start < n:
        output = update_tile(tail_start, n - tail_start, output)
    return output


def _pair_mlp_parts_tiled(mlp, parts, *, tile_size: int):

    parts = tuple(parts)
    if not parts:
        raise ValueError("pair MLP requires at least one input part")
    pair_shape = parts[0].shape[:2]
    if any(part.shape[:2] != pair_shape for part in parts):
        raise ValueError("pair MLP input parts must share their [R,N] axes")
    n = pair_shape[1]
    tile_width = min(n, int(tile_size))
    if tile_width < 1:
        raise ValueError("pair MLP tile_size must be positive")
    full_tiles = n // tile_width
    tail_start = full_tiles * tile_width
    output = jnp.zeros(
        pair_shape + (int(mlp.out_proj.weight.shape[1]),),
        dtype=parts[0].dtype,
    )

    def update_tile(start, width, current):
        tile = jnp.concatenate(
            tuple(
                jax.lax.dynamic_slice_in_dim(part, start, width, axis=1)
                for part in parts
            ),
            axis=-1,
        )
        return jax.lax.dynamic_update_slice_in_dim(
            current, _mlp(mlp, tile), start, axis=1
        )

    output = jax.lax.fori_loop(
        0,
        full_tiles,
        lambda tile_index, current: update_tile(
            tile_index * tile_width, tile_width, current
        ),
        output,
    )
    if tail_start < n:
        output = update_tile(tail_start, n - tail_start, output)
    return output


def _sequence_context_attention(
    layer,
    c_rows,
    edge_rows,
    bmask,
    row_indices,
    *,
    axis_name: str,
    block_k: int,
    tile_size: int,
):

    r = c_rows.shape[0]
    n = bmask.shape[0]
    qkv_rows = _linear(layer.W_QKV, c_rows).reshape(
        r, 3, layer.n_heads_kernel, layer.d_head
    )
    qkv_all = _gather_rows(qkv_rows, axis_name=axis_name)
    query = qkv_rows[:, 0]
    key = qkv_all[:, 1]
    value = qkv_all[:, 2]
    col_indices = jnp.arange(n, dtype=jnp.int32)
    rel = row_indices[:, None] - col_indices[None, :]
    if isinstance(layer, PhysicalReadoutContextLayer):
        direction = jnp.where(rel < 0, 1.0, jnp.where(rel > 0, -1.0, 0.0)).astype(
            edge_rows.dtype
        )[..., None]
    elif isinstance(layer, RouterContextLayer):
        direction = jnp.zeros((r, n, 1), dtype=edge_rows.dtype)
    else:
        raise TypeError("unsupported contextualizer layer")
    bias = _pair_mlp_parts_tiled(
        layer.bias_mlp,
        (edge_rows, direction),
        tile_size=tile_size,
    )
    bias = bias / jnp.sqrt(jnp.asarray(layer.d_head, dtype=bias.dtype))
    if isinstance(layer, PhysicalReadoutContextLayer):
        bias = bias + jnp.transpose(
            lca_alibi_bias(
                row_indices,
                col_indices,
                lca_fixed_slopes(layer.n_heads_kernel, dtype=bias.dtype),
            ),
            (1, 2, 0),
        )
    out = pallas_rectangular_edge_attention(
        query,
        key,
        value,
        bias,
        bmask,
        block_k=block_k,
    )
    out = jax.nn.sigmoid(out[:, : layer.n_heads]) * out[:, layer.n_heads :]
    return _linear(layer.W_O, out.reshape(r, -1))


def _sequence_context_layer(
    layer,
    c_rows,
    edge_rows,
    g,
    bmask,
    *,
    axis_name: str,
    axis_size: int,
    tile_size: int,
    attention_block_k: int,
):

    r, n = edge_rows.shape[:2]
    dtype = edge_rows.dtype
    row_indices = _global_row_indices(axis_name=axis_name, local_size=r)
    row_mask = bmask[row_indices]
    edge_n = _norm(layer.ln_edge, edge_rows)
    if isinstance(layer, PhysicalReadoutContextLayer):
        reverse_n = transpose_pair_rows_local(
            edge_n,
            axis_name=axis_name,
            axis_size=axis_size,
        )
        clock_rows = _local_rows(
            layer._slot_clock(n, dtype, bmask),
            axis_name=axis_name,
            local_size=r,
        )
        summary = layer.edge_summary_tiled(
            edge_n,
            bmask.astype(dtype),
            edge_reverse_rows=reverse_n,
            row_indices=row_indices,
            tile_size=tile_size,
        )
        parts = [
            _norm(layer.ln_c, c_rows + clock_rows),
            _norm(layer.ln_summary, summary),
        ]
    elif isinstance(layer, RouterContextLayer):
        clock_rows = None
        parts = [_norm(layer.ln_c, c_rows)]
    else:
        raise TypeError("unsupported contextualizer layer")
    delta_ctx = layer.residual_scale * _mlp(
        layer.ctx_mlp, jnp.concatenate(parts, axis=-1)
    )
    c1 = jnp.where(
        row_mask[:, None].astype(bool),
        c_rows + delta_ctx,
        jnp.zeros_like(c_rows),
    )

    c1_all = _gather_rows(c1, axis_name=axis_name)
    g = _g_update(
        layer.g_update,
        g,
        _g_descriptor_pool(layer.g_pool, g, c1_all, bmask),
    )

    edge_ctx_rows = _linear(
        layer.edge_node_ctx_proj,
        _norm(layer.ln_edge_ctx, c1),
    )
    edge_ctx_all = _gather_rows(edge_ctx_rows, axis_name=axis_name)
    edge1 = layer.edge_update_tiled(
        edge_rows,
        edge_n,
        edge_ctx_rows,
        edge_ctx_all,
        bmask.astype(dtype),
        row_indices=row_indices,
        g=g,
        tile_size=tile_size,
    )

    attention_source = c1
    if clock_rows is not None:
        attention_source = attention_source + clock_rows
    delta_attn = layer.residual_scale * _sequence_context_attention(
        layer,
        _norm(layer.ln_attn, attention_source),
        _norm(layer.ln_edge_attn, edge1),
        bmask,
        row_indices,
        axis_name=axis_name,
        block_k=attention_block_k,
        tile_size=tile_size,
    )
    c_out = jnp.where(
        row_mask[:, None].astype(bool),
        c1 + delta_attn,
        jnp.zeros_like(c1),
    )
    return c_out, edge1, g


def sequence_parallel_contextualizer_local(
    contextualizer,
    node_rows,
    edge_rows,
    real_mask,
    structural_mask,
    g=None,
    *,
    axis_name: str,
    axis_size: int,
    tile_size: int = 128,
    attention_block_k: int = 128,
):

    if not isinstance(contextualizer, (PhysicalReadoutContext, RouterContext)):
        raise TypeError("unsupported contextualizer")
    node_rows = node_rows.astype(jnp.float32)
    edge_rows = edge_rows.astype(jnp.float32)
    r, n = edge_rows.shape[:2]
    if node_rows.shape[0] != r or n != r * axis_size:
        raise ValueError("contextualizer inputs do not match the seq row layout")
    if real_mask.shape != (n,) or structural_mask.shape != (n,):
        raise ValueError("contextualizer masks must have replicated shape [N]")
    row_indices = _global_row_indices(axis_name=axis_name, local_size=r)
    real_rows = real_mask[row_indices].astype(bool)
    active_rows = structural_mask[row_indices].astype(bool)
    virtual_rows = active_rows & ~real_rows
    real = real_mask.astype(bool)
    active = structural_mask.astype(bool)
    virtual = active & ~real

    virtual_node = contextualizer.virtual_node[0].astype(node_rows.dtype)
    node_rows = jnp.where(
        real_rows[:, None],
        node_rows,
        jnp.where(
            virtual_rows[:, None],
            virtual_node[None, :],
            jnp.zeros_like(node_rows),
        ),
    )
    row_real_pair = real_rows[:, None] & real[None, :]
    mixed_pair = (real_rows[:, None] & virtual[None, :]) | (
        virtual_rows[:, None] & real[None, :]
    )
    virtual_pair = virtual_rows[:, None] & virtual[None, :]
    edge_rows = jnp.where(
        row_real_pair[..., None],
        edge_rows,
        jnp.where(
            mixed_pair[..., None],
            contextualizer.edge_empty_nonempty[0].astype(edge_rows.dtype),
            jnp.where(
                virtual_pair[..., None],
                contextualizer.edge_empty_empty[0].astype(edge_rows.dtype),
                jnp.zeros_like(edge_rows),
            ),
        ),
    )

    step = partial(
        _sequence_context_layer,
        bmask=structural_mask,
        axis_name=axis_name,
        axis_size=axis_size,
        tile_size=tile_size,
        attention_block_k=attention_block_k,
    )
    dynamic, static = eqx.partition(contextualizer.layers, eqx.is_array)

    def scan_step(carry, layer_dynamic):
        layer = eqx.combine(layer_dynamic, static)
        return step(layer, *carry), None

    (node_rows, edge_rows, g), _ = jax.lax.scan(
        scan_step, (node_rows, edge_rows, g), dynamic
    )
    return node_rows, edge_rows, g


def _sequence_tree_fwl(
    module,
    edge_rows,
    c_rows,
    c_all,
    mask,
    row_mask,
    *,
    axis_name: str,
    axis_size: int,
    tile_size: int,
):

    width = int(edge_rows.shape[1])
    if tile_size < 1:
        raise ValueError("tree FWL tile_size must be positive")

    tile_width = gcd(width, min(width, int(tile_size)))
    n_tiles = width // tile_width
    c_rows_ctx = _norm(module.ln_c, c_rows)
    c_all_ctx = _norm(module.ln_c, c_all)
    if module.node_ctx_proj is not None:
        c_rows_ctx = _linear(module.node_ctx_proj, c_rows_ctx)
        c_all_ctx = _linear(module.node_ctx_proj, c_all_ctx)

    left0 = jnp.zeros(
        (
            edge_rows.shape[0],
            width,
            int(module.psi_L_out.weight.shape[1]),
        ),
        dtype=edge_rows.dtype,
    )

    def project_left_tile(tile_index, left):
        start = tile_index * tile_width
        edge_tile = jax.lax.dynamic_slice_in_dim(edge_rows, start, tile_width, axis=1)
        c_columns = jax.lax.dynamic_slice_in_dim(c_all_ctx, start, tile_width, axis=0)
        mask_columns = jax.lax.dynamic_slice_in_dim(mask, start, tile_width, axis=0)
        pair_parts = (
            _norm(module.ln_edge, edge_tile),
            c_rows_ctx[:, None, :],
            c_columns[None, :, :],
        )
        projected = _psi_project(module, pair_parts, left=True)
        projected = projected * (
            row_mask[:, None, None].astype(projected.dtype)
            * mask_columns[None, :, None].astype(projected.dtype)
        )
        return jax.lax.dynamic_update_slice_in_dim(left, projected, start, axis=1)

    left = jax.lax.fori_loop(0, n_tiles, project_left_tile, left0)
    n_eff = jnp.maximum(jnp.sum(mask), 1.0).astype(edge_rows.dtype)

    def update_destination_tile(tile_index, updated_edges):
        start = tile_index * tile_width
        edge_tile = jax.lax.dynamic_slice_in_dim(
            updated_edges, start, tile_width, axis=1
        )
        c_columns = jax.lax.dynamic_slice_in_dim(c_all_ctx, start, tile_width, axis=0)
        mask_columns = jax.lax.dynamic_slice_in_dim(mask, start, tile_width, axis=0)
        pair_parts = (
            _norm(module.ln_edge, edge_tile),
            c_rows_ctx[:, None, :],
            c_columns[None, :, :],
        )
        right = _psi_project(module, pair_parts, left=False)
        right = right * mask_columns[None, :, None].astype(right.dtype)
        path = ring_learned_fwl2_columns_local(
            left,
            right,
            axis_name=axis_name,
            axis_size=axis_size,
        )
        path = path / jnp.sqrt(n_eff)
        path = _norm(module.ln_path, path)
        hidden = fused_silu(_split_linear(module.ffn_in, (*pair_parts, path)))
        delta = _linear(module.ffn_out, hidden)
        update_mask = row_mask[:, None].astype(bool) & mask_columns[None, :].astype(
            bool
        )
        edge_tile = _tree_ngpt_residual(
            edge_tile,
            delta,
            module.alpha,
            max_gain=module.ngpt_alpha_max,
            tag_id="",
            update_mask=update_mask,
        )
        return jax.lax.dynamic_update_slice_in_dim(
            updated_edges, edge_tile, start, axis=1
        )

    return jax.lax.fori_loop(0, n_tiles, update_destination_tile, edge_rows)


def _sequence_level_edge_attention(
    module,
    c_rows,
    edge_rows,
    mask,
    row_mask,
    *,
    axis_name: str,
    level: int,
    tile_size: int,
    attention_block_k: int,
):

    r, n = edge_rows.shape[:2]
    x = _raw_norm(
        c_rows,
        module.ln_scale,
        None,
        eps=module.ln_eps,
    )
    qkv_rows = (x @ module.w_qkv).reshape(r, 3, module.n_heads_kernel, module.d_head)
    qkv_all = _gather_rows(qkv_rows, axis_name=axis_name)
    query = qkv_rows[:, 0]
    key = qkv_all[:, 1]
    value = qkv_all[:, 2]
    bias = _pair_mlp_tiled(module.bias_mlp, edge_rows, tile_size=tile_size)
    bias = bias / jnp.sqrt(jnp.asarray(module.d_head, dtype=bias.dtype))
    row_indices = _global_row_indices(axis_name=axis_name, local_size=r)
    bias = bias + jnp.transpose(
        lca_alibi_bias(
            row_indices,
            jnp.arange(n, dtype=jnp.int32),
            lca_fixed_slopes(module.n_heads_kernel, dtype=bias.dtype),
        ),
        (1, 2, 0),
    )
    out = pallas_rectangular_edge_attention(
        query,
        key,
        value,
        bias,
        mask,
        block_k=attention_block_k,
    )
    out = jax.nn.sigmoid(out[:, : module.n_heads]) * out[:, module.n_heads :]
    proposal_attn = row_mask[:, None].astype(out.dtype) * (
        out.reshape(r, -1) @ module.w_o
    )
    c_attn = _tree_ngpt_residual(
        c_rows,
        proposal_attn,
        module.alpha_attn,
        max_gain=module.ngpt_alpha_max,
        tag_id="",
        update_mask=row_mask,
    )
    x_ffn = _raw_norm(
        c_attn,
        module.ffn_ln_scale,
        None,
        eps=module.ln_eps,
    )
    proposal_ffn = row_mask[:, None].astype(x_ffn.dtype) * (
        fused_silu(x_ffn @ module.ffn_w1 + module.ffn_b1) @ module.ffn_w2
        + module.ffn_b2
    )
    return _tree_ngpt_residual(
        c_attn,
        proposal_ffn,
        module.alpha_ffn,
        max_gain=module.ngpt_alpha_max,
        tag_id="",
        update_mask=row_mask,
    )


def sequence_parallel_physical_leaf_local(
    kernel,
    contextualized_node_rows,
    global_stream,
    *,
    axis_name: str,
):

    from hamiltonzero.compiled.tree import _project_global, compile_target_leaf_h

    leaf_g_emb = _project_global(
        kernel.leaf_projection,
        global_stream,
        dense_tag="gladder.to_gemb",
        norm_tag="gladder.gemb_ln",
    )
    node_all = _gather_rows(contextualized_node_rows, axis_name=axis_name)
    leaf_h = compile_target_leaf_h(kernel.leaf, node_all, leaf_g_emb)
    c_rows = _linear(kernel.leaf.P_c, contextualized_node_rows)
    c_rows = _tree_sphere(c_rows)
    return leaf_h, c_rows


def sequence_parallel_reduce_physical_local(
    kernel,
    edge_rows,
    leaf_h,
    c_rows,
    leaf_real,
    structural_mask,
    global_stream,
    permutation,
    *,
    axis_name: str,
    axis_size: int,
    replicate_threshold: int = 512,
    contextualizer_tile_size: int = 128,
    attention_block_k: int = 128,
):

    from hamiltonzero.compiled.tree import (
        compile_merge_h,
        compile_physical_tree_from_reduced_state,
    )
    from hamiltonzero.compiled.types import CARRY_LEFT, CARRY_RIGHT, EMPTY, MERGE
    from hamiltonzero.model.tree import (
        _tree_active_clock_depth,
        _tree_depth_count_features,
        edge_merge_masked,
    )

    if replicate_threshold < 1:
        raise ValueError("replicate_threshold must be positive")
    local_size, n = edge_rows.shape[:2]
    if n != local_size * axis_size or n & (n - 1):
        raise ValueError("physical sequence compiler requires power-of-two N")
    if local_size & (local_size - 1):
        raise ValueError("each seq lane must own a power-of-two row count")

    g = global_stream

    merge = kernel.merge
    m = leaf_real.astype(c_rows.dtype)
    k = structural_mask.astype(c_rows.dtype)
    counts = m
    n_total = jnp.sum(m)
    feature_n_levels = _tree_active_clock_depth(m)
    clock_depth = _tree_active_clock_depth(k)
    edge_rows = _tree_sphere(edge_rows)

    early_merge_h = []
    early_opcodes = []
    width = n
    rows_per_lane = local_size
    level = 0
    while width > int(replicate_threshold):
        if rows_per_lane < 2:
            raise ValueError(
                "replicate_threshold is too small for the available seq lanes"
            )
        c_all = _gather_rows(c_rows, axis_name=axis_name)
        c_a_rows, c_b_rows = c_rows[0::2], c_rows[1::2]
        m_a, m_b = m[0::2], m[1::2]
        k_a, k_b = k[0::2], k[1::2]
        cnt_a, cnt_b = counts[0::2], counts[1::2]
        m_rows = _local_rows(m, axis_name=axis_name, local_size=rows_per_lane)
        k_rows = _local_rows(k, axis_name=axis_name, local_size=rows_per_lane)
        m_a_rows, m_b_rows = m_rows[0::2], m_rows[1::2]
        k_a_rows, k_b_rows = k_rows[0::2], k_rows[1::2]
        both_struct = k_a * k_b
        both_struct_rows = _local_rows(
            both_struct,
            axis_name=axis_name,
            local_size=rows_per_lane // 2,
        )
        pair_base = jnp.maximum(
            jnp.sum((k_a + k_b - k_a * k_b).astype(jnp.int32)),
            jnp.asarray(2, dtype=jnp.int32),
        )
        depth = _tree_depth_count_features(
            cnt_a,
            cnt_b,
            n_total,
            level,
            feature_n_levels,
            c_rows.dtype,
        )
        depth_rows = _local_rows(
            depth,
            axis_name=axis_name,
            local_size=rows_per_lane // 2,
        )

        edge_blocks = edge_rows.reshape(
            rows_per_lane // 2,
            2,
            width // 2,
            2,
            edge_rows.shape[-1],
        )
        local_parent = jnp.arange(rows_per_lane // 2, dtype=jnp.int32)
        global_parent = (
            jax.lax.axis_index(axis_name) * (rows_per_lane // 2) + local_parent
        )
        sibling_lr = edge_blocks[local_parent, 0, global_parent, 1]
        sibling_rl = edge_blocks[local_parent, 1, global_parent, 0]
        level_active = jnp.any(both_struct.astype(bool))
        g_level = g @ kernel.tree_projection_weight + kernel.tree_projection_bias

        def candidate_one(ca, cb, elr, erl, dep, pidx, active):
            return merge.context_candidate(
                ca,
                cb,
                g_level,
                sibling_edge_lr=elr,
                sibling_edge_rl=erl,
                level_idx=jnp.int32(level),
                pair_idx=pidx,
                pair_base=pair_base,
                clock_depth=clock_depth,
                depth_feats=dep,
                kfac_structural_mask=active,
                kfac_g_structural_mask=level_active,
                kfac_scan_shared=False,
            )

        candidate = jax.vmap(candidate_one)(
            c_a_rows,
            c_b_rows,
            sibling_lr,
            sibling_rl,
            depth_rows,
            global_parent,
            both_struct_rows,
        )
        early_merge_h.append(
            _gather_rows(
                compile_merge_h(
                    merge,
                    candidate,
                    depth_rows,
                ),
                axis_name=axis_name,
            )
        )
        early_opcodes.append(
            jnp.where(
                m_a.astype(bool),
                jnp.where(m_b.astype(bool), MERGE, CARRY_LEFT),
                jnp.where(m_b.astype(bool), CARRY_RIGHT, EMPTY),
            ).astype(jnp.uint8)
        )

        gate_a_mask, gate_b_mask = k_a_rows, k_b_rows
        gate_both = gate_a_mask * gate_b_mask
        c_rows = (
            gate_both[:, None] * candidate
            + (gate_a_mask * (1.0 - gate_b_mask))[:, None] * c_a_rows
            + ((1.0 - gate_a_mask) * gate_b_mask)[:, None] * c_b_rows
        )

        c_a_all, c_b_all = c_all[0::2], c_all[1::2]

        parent_width = width // 2
        edge_tile_width = gcd(
            parent_width,
            min(parent_width, int(contextualizer_tile_size)),
        )
        edge_tile_count = parent_width // edge_tile_width
        edge_new0 = jnp.zeros(
            (rows_per_lane // 2, parent_width, edge_rows.shape[-1]),
            dtype=edge_rows.dtype,
        )

        def merge_destination_tile(tile_index, edge_output):
            start = tile_index * edge_tile_width

            def take_columns(x):
                return jax.lax.dynamic_slice_in_dim(x, start, edge_tile_width, axis=0)

            m_a_tile, m_b_tile = take_columns(m_a), take_columns(m_b)
            k_a_tile, k_b_tile = take_columns(k_a), take_columns(k_b)
            c_a_tile, c_b_tile = take_columns(c_a_all), take_columns(c_b_all)

            edge_block_tile = jax.lax.dynamic_slice_in_dim(
                edge_blocks, start, edge_tile_width, axis=2
            )

            def edge_row(e0, e1, e2, e3, ma, mb, ka, kb, ca, cb):
                return jax.vmap(
                    lambda x0, x1, x2, x3, mqa, mqb, kqa, kqb, cqa, cqb: (
                        edge_merge_masked(
                            x0,
                            x1,
                            x2,
                            x3,
                            ma,
                            mb,
                            mqa,
                            mqb,
                            ca,
                            cb,
                            cqa,
                            cqb,
                            merge.edge_merge,
                            k_2i=ka,
                            k_2i1=kb,
                            k_2j=kqa,
                            k_2j1=kqb,
                            kfac_scan_shared=False,
                        )[0]
                    )
                )(
                    e0,
                    e1,
                    e2,
                    e3,
                    m_a_tile,
                    m_b_tile,
                    k_a_tile,
                    k_b_tile,
                    c_a_tile,
                    c_b_tile,
                )

            edge_tile = jax.vmap(edge_row)(
                edge_block_tile[:, 0, :, 0],
                edge_block_tile[:, 0, :, 1],
                edge_block_tile[:, 1, :, 0],
                edge_block_tile[:, 1, :, 1],
                m_a_rows,
                m_b_rows,
                k_a_rows,
                k_b_rows,
                c_a_rows,
                c_b_rows,
            )
            return jax.lax.dynamic_update_slice_in_dim(
                edge_output, edge_tile, start, axis=1
            )

        edge_new = jax.lax.fori_loop(
            0, edge_tile_count, merge_destination_tile, edge_new0
        )
        c_all_new = _gather_rows(c_rows, axis_name=axis_name)
        edge_new = _sequence_tree_fwl(
            merge.tree_edge_fwl,
            edge_new,
            c_rows,
            c_all_new,
            both_struct,
            both_struct_rows,
            axis_name=axis_name,
            axis_size=axis_size,
            tile_size=contextualizer_tile_size,
        )
        edge_keep = both_struct_rows[:, None] * both_struct[None, :]

        def finalize_edge_tile(tile_index, current_edges):
            start = tile_index * edge_tile_width
            current_tile = jax.lax.dynamic_slice_in_dim(
                current_edges, start, edge_tile_width, axis=1
            )
            fallback_tile = jax.lax.dynamic_slice_in_dim(
                edge_blocks, start, edge_tile_width, axis=2
            )[:, 0, :, 0]
            keep_tile = jax.lax.dynamic_slice_in_dim(
                edge_keep, start, edge_tile_width, axis=1
            )[..., None].astype(bool)
            current_tile = _tree_sphere(current_tile)
            current_tile = jnp.where(keep_tile, current_tile, fallback_tile)
            return jax.lax.dynamic_update_slice_in_dim(
                current_edges, current_tile, start, axis=1
            )

        edge_rows = jax.lax.fori_loop(0, edge_tile_count, finalize_edge_tile, edge_new)
        c_skip = c_rows
        c_rows = _sequence_level_edge_attention(
            merge.level_edge_attn,
            c_rows,
            edge_rows,
            both_struct,
            both_struct_rows,
            axis_name=axis_name,
            level=level,
            tile_size=contextualizer_tile_size,
            attention_block_k=attention_block_k,
        )
        c_rows = jnp.where(
            both_struct_rows[:, None].astype(bool),
            _tree_sphere(c_rows),
            c_skip,
        )

        m = m_a + m_b - m_a * m_b
        k = k_a + k_b - k_a * k_b
        counts = cnt_a + cnt_b
        level_mask = k
        c_all_new = _gather_rows(c_rows, axis_name=axis_name)
        updated_g = _g_update(
            kernel.tree_update,
            g,
            _g_descriptor_pool(kernel.tree_pool, g, c_all_new, level_mask),
        )
        g = jnp.where(level_active, updated_g, g)
        width //= 2
        rows_per_lane //= 2
        level += 1

    c_reduced = _gather_rows(c_rows, axis_name=axis_name)
    edge_reduced = _gather_rows(edge_rows, axis_name=axis_name)
    return compile_physical_tree_from_reduced_state(
        kernel,
        perm=permutation,
        leaf_real=leaf_real,
        leaf_h=leaf_h,
        c_reduced=c_reduced,
        edge_reduced=edge_reduced,
        real_reduced=m,
        structural_reduced=k,
        counts_reduced=counts,
        g_reduced=g,
        early_merge_h=tuple(early_merge_h),
        early_opcodes=tuple(early_opcodes),
        full_structural_mask=structural_mask,
    )


def sequence_parallel_shared_trunk_local(
    kernel: TrunkCompilerKernel,
    j_double_prime_rows: jax.Array,
    h_prime: jax.Array,
    real_mask: jax.Array,
    balanced_mask: jax.Array,
    *,
    axis_name: str,
    axis_size: int,
    featurizer_tile_size: int = 128,
    attention_block_k: int = 128,
) -> SharedTrunk:

    if j_double_prime_rows.ndim != 3 or j_double_prime_rows.shape[-1] != 10:
        raise ValueError("J rows must have shape [N/P,N,10]")
    local_size, n = j_double_prime_rows.shape[:2]
    if n != local_size * axis_size:
        raise ValueError(
            f"sequence rows must evenly tile N: local={local_size}, "
            f"N={n}, shards={axis_size}"
        )
    if h_prime.shape != (n, 3):
        raise ValueError(f"h_prime must have shape {(n, 3)}, got {h_prime.shape}")
    if real_mask.shape != (n,) or balanced_mask.shape != (n,):
        raise ValueError("real and balanced masks must both have global shape [N]")
    if j_double_prime_rows.dtype != jnp.float32:
        raise TypeError("fresh sequence compiler currently requires fp32 graph inputs")

    row_indices = _global_row_indices(axis_name=axis_name, local_size=local_size)
    row_mask = real_mask[row_indices]
    featurizer = kernel.featurizer
    bond_rows, descriptor_rows = featurizer.eval_embed_local_rows(
        j_double_prime_rows,
        row_mask,
        real_mask,
        tile_size=featurizer_tile_size,
    )
    descriptor_all = _gather_rows(descriptor_rows, axis_name=axis_name)

    sum_j2, count_j = featurizer.eval_jh_stats_rows(
        j_double_prime_rows,
        row_mask,
        real_mask,
        row_indices=row_indices,
    )
    jh_stats = (
        jax.lax.psum(sum_j2, axis_name),
        jax.lax.psum(count_j, axis_name),
    )

    local_rows, global_raw = featurizer.eval_finalize_local_rows(
        J_double_prime_rows=j_double_prime_rows,
        local_desc_rows=descriptor_rows,
        local_desc_all=descriptor_all,
        row_indices=row_indices,
        mask=real_mask,
        h_prime=h_prime,
        jh_stats=jh_stats,
    )
    local_all = _gather_rows(local_rows, axis_name=axis_name)
    edge_rows = featurizer.eval_edge_rows(
        bond_emb_rows=bond_rows,
        local_rows=local_rows,
        local_final_all=local_all,
        global_feat=global_raw,
        row_indices=row_indices,
        mask=real_mask,
        tile_size=featurizer_tile_size,
    )

    g = _tree_sphere(global_raw.astype(local_rows.dtype))
    local_rows, edge_rows, g = _sequence_trunk_local(
        kernel.trunk,
        local_rows,
        edge_rows,
        g,
        real_mask,
        row_mask,
        axis_name=axis_name,
        axis_size=axis_size,
        attention_block_k=attention_block_k,
    )
    global_stream = _edge_row_col_global_update(
        kernel.shared_global,
        g,
        edge_rows,
        real_mask,
        row_mask,
        axis_name=axis_name,
    )
    return SharedTrunk(
        node_raw=local_rows,
        edge_raw=edge_rows,
        global_raw=global_raw,
        global_stream=global_stream,
        real_mask=real_mask,
        balanced_mask=balanced_mask,
    )


def build_sequence_parallel_shared_trunk(
    *,
    mesh: Mesh,
    kernel_template: TrunkCompilerKernel,
    axis_name: str = "seq",
    featurizer_tile_size: int = 128,
    attention_block_k: int = 128,
) -> Callable[
    [TrunkCompilerKernel, jax.Array, jax.Array, jax.Array, jax.Array],
    SharedTrunk,
]:

    if tuple(mesh.axis_names) != (axis_name,):
        raise ValueError(
            f"sequence trunk requires a one-dimensional {axis_name!r} mesh; "
            f"got {mesh.axis_names}"
        )
    axis_size = int(mesh.shape[axis_name])
    replicated_spec = P()
    edge_spec = P(axis_name, None, None)
    kernel_specs = jax.tree_util.tree_map(lambda _: replicated_spec, kernel_template)
    output_specs = SharedTrunk(
        node_raw=P(axis_name, None),
        edge_raw=edge_spec,
        global_raw=replicated_spec,
        global_stream=replicated_spec,
        real_mask=replicated_spec,
        balanced_mask=replicated_spec,
    )

    local = partial(
        sequence_parallel_shared_trunk_local,
        axis_name=axis_name,
        axis_size=axis_size,
        featurizer_tile_size=featurizer_tile_size,
        attention_block_k=attention_block_k,
    )
    mapped = jax.shard_map(
        local,
        mesh=mesh,
        in_specs=(
            kernel_specs,
            edge_spec,
            replicated_spec,
            replicated_spec,
            replicated_spec,
        ),
        out_specs=output_specs,
        check_vma=False,
    )

    replicated = NamedSharding(mesh, replicated_spec)
    edge_sharding = NamedSharding(mesh, edge_spec)
    output_shardings = SharedTrunk(
        node_raw=NamedSharding(mesh, P(axis_name, None)),
        edge_raw=edge_sharding,
        global_raw=replicated,
        global_stream=replicated,
        real_mask=replicated,
        balanced_mask=replicated,
    )
    return jax.jit(
        mapped,
        in_shardings=(
            jax.tree_util.tree_map(lambda _: replicated, kernel_template),
            edge_sharding,
            replicated,
            replicated,
            replicated,
        ),
        out_shardings=output_shardings,
    )


def build_sequence_parallel_contextualizer(
    *,
    mesh: Mesh,
    contextualizer_template,
    g_template: jax.Array,
    axis_name: str = "seq",
    tile_size: int = 128,
    attention_block_k: int = 128,
) -> Callable:

    if tuple(mesh.axis_names) != (axis_name,):
        raise ValueError(
            f"contextualizer requires a one-dimensional {axis_name!r} mesh"
        )
    axis_size = int(mesh.shape[axis_name])
    rep_spec = P()
    node_spec = P(axis_name, None)
    edge_spec = P(axis_name, None, None)
    context_specs = jax.tree_util.tree_map(lambda _: rep_spec, contextualizer_template)
    local = partial(
        sequence_parallel_contextualizer_local,
        axis_name=axis_name,
        axis_size=axis_size,
        tile_size=tile_size,
        attention_block_k=attention_block_k,
    )
    mapped = jax.shard_map(
        local,
        mesh=mesh,
        in_specs=(
            context_specs,
            node_spec,
            edge_spec,
            rep_spec,
            rep_spec,
            rep_spec,
        ),
        out_specs=(node_spec, edge_spec, rep_spec),
        check_vma=False,
    )
    rep = NamedSharding(mesh, rep_spec)
    node_sharding = NamedSharding(mesh, node_spec)
    edge_sharding = NamedSharding(mesh, edge_spec)
    return jax.jit(
        mapped,
        in_shardings=(
            jax.tree_util.tree_map(lambda _: rep, contextualizer_template),
            node_sharding,
            edge_sharding,
            rep,
            rep,
            rep,
        ),
        out_shardings=(node_sharding, edge_sharding, rep),
    )


def build_sequence_pair_permute(*, mesh: Mesh, axis_name: str = "seq") -> Callable:

    if tuple(mesh.axis_names) != (axis_name,):
        raise ValueError(
            f"pair permutation requires a one-dimensional {axis_name!r} mesh"
        )
    axis_size = int(mesh.shape[axis_name])
    node_spec = P(axis_name, None)
    edge_spec = P(axis_name, None, None)

    def local(node_rows, edge_rows, permutation):
        return (
            ring_permute_rows_local(
                node_rows,
                permutation,
                axis_name=axis_name,
                axis_size=axis_size,
            ),
            permute_pair_rows_and_columns_local(
                edge_rows,
                permutation,
                axis_name=axis_name,
                axis_size=axis_size,
            ),
        )

    return jax.jit(
        jax.shard_map(
            local,
            mesh=mesh,
            in_specs=(node_spec, edge_spec, P()),
            out_specs=(node_spec, edge_spec),
            check_vma=False,
        )
    )


def build_sequence_parallel_edge_global_update(
    *,
    mesh: Mesh,
    module_template,
    axis_name: str = "seq",
    tile_size: int = 128,
) -> Callable:

    if tuple(mesh.axis_names) != (axis_name,):
        raise ValueError(
            f"global edge update requires a one-dimensional {axis_name!r} mesh"
        )
    rep_spec = P()
    edge_spec = P(axis_name, None, None)
    mapped = jax.shard_map(
        partial(
            sequence_parallel_edge_global_update_local,
            axis_name=axis_name,
            tile_size=tile_size,
        ),
        mesh=mesh,
        in_specs=(
            jax.tree_util.tree_map(lambda _: rep_spec, module_template),
            rep_spec,
            edge_spec,
            rep_spec,
        ),
        out_specs=rep_spec,
        check_vma=False,
    )
    rep = NamedSharding(mesh, rep_spec)
    return jax.jit(
        mapped,
        in_shardings=(
            jax.tree_util.tree_map(lambda _: rep, module_template),
            rep,
            NamedSharding(mesh, edge_spec),
            rep,
        ),
        out_shardings=rep,
    )


def build_sequence_parallel_physical_leaf(
    *,
    mesh: Mesh,
    kernel_template,
    axis_name: str = "seq",
) -> Callable:

    if tuple(mesh.axis_names) != (axis_name,):
        raise ValueError(
            f"physical leaf projection requires a one-dimensional {axis_name!r} mesh"
        )
    rep_spec = P()
    node_spec = P(axis_name, None)
    kernel_specs = jax.tree_util.tree_map(lambda _: rep_spec, kernel_template)
    mapped = jax.shard_map(
        partial(sequence_parallel_physical_leaf_local, axis_name=axis_name),
        mesh=mesh,
        in_specs=(kernel_specs, node_spec, rep_spec),
        out_specs=((rep_spec,), node_spec),
        check_vma=False,
    )
    rep = NamedSharding(mesh, rep_spec)
    node_sharding = NamedSharding(mesh, node_spec)
    return jax.jit(
        mapped,
        in_shardings=(
            jax.tree_util.tree_map(lambda _: rep, kernel_template),
            node_sharding,
            rep,
        ),
        out_shardings=((rep,), node_sharding),
    )


def build_sequence_parallel_physical_reducer(
    *,
    mesh: Mesh,
    kernel_template,
    edge_template,
    axis_name: str = "seq",
    replicate_threshold: int = 512,
    contextualizer_tile_size: int = 128,
    attention_block_k: int = 128,
) -> Callable:

    from hamiltonzero.compiled.types import CompiledTree

    if tuple(mesh.axis_names) != (axis_name,):
        raise ValueError(
            f"physical reducer requires a one-dimensional {axis_name!r} mesh"
        )
    n = int(edge_template.shape[0])
    if tuple(edge_template.shape[:2]) != (n, n):
        raise ValueError("edge_template must have square global pair axes")
    axis_size = int(mesh.shape[axis_name])
    if n % axis_size or n & (n - 1):
        raise ValueError("global N must be a power of two divisible by seq lanes")
    rep_spec = P()
    node_spec = P(axis_name, None)
    edge_spec = P(axis_name, None, None)
    kernel_specs = jax.tree_util.tree_map(lambda _: rep_spec, kernel_template)
    n_levels = n.bit_length() - 1
    output_specs = CompiledTree(
        perm=rep_spec,
        inv_perm=rep_spec,
        leaf_real=rep_spec,
        leaf_h=(rep_spec,),
        leaf_combiner_h=(),
        merge_h=(rep_spec,) * n_levels,
        opcodes=(rep_spec,) * n_levels,
        readout_h=(rep_spec,),
        readout_combiner_h=(),
    )
    local = partial(
        sequence_parallel_reduce_physical_local,
        axis_name=axis_name,
        axis_size=axis_size,
        replicate_threshold=replicate_threshold,
        contextualizer_tile_size=contextualizer_tile_size,
        attention_block_k=attention_block_k,
    )
    mapped = jax.shard_map(
        local,
        mesh=mesh,
        in_specs=(
            kernel_specs,
            edge_spec,
            (rep_spec,),
            node_spec,
            rep_spec,
            rep_spec,
            rep_spec,
            rep_spec,
        ),
        out_specs=output_specs,
        check_vma=False,
    )
    rep = NamedSharding(mesh, rep_spec)
    node_sharding = NamedSharding(mesh, node_spec)
    edge_sharding = NamedSharding(mesh, edge_spec)
    output_shardings = jax.tree_util.tree_map(lambda _: rep, output_specs)
    return jax.jit(
        mapped,
        in_shardings=(
            jax.tree_util.tree_map(lambda _: rep, kernel_template),
            edge_sharding,
            (rep,),
            node_sharding,
            rep,
            rep,
            rep,
            rep,
        ),
        out_shardings=output_shardings,
    )


__all__ = [
    "build_sequence_pair_permute",
    "build_sequence_parallel_contextualizer",
    "build_sequence_parallel_edge_global_update",
    "build_sequence_parallel_physical_leaf",
    "build_sequence_parallel_physical_reducer",
    "build_sequence_parallel_shared_trunk",
    "permute_pair_rows_and_columns_local",
    "ring_permute_rows_local",
    "sequence_parallel_contextualizer_local",
    "sequence_parallel_edge_global_update_local",
    "sequence_parallel_physical_leaf_local",
    "sequence_parallel_reduce_physical_local",
    "sequence_parallel_shared_trunk_local",
    "transpose_pair_rows_local",
]