src/hamiltonzero/model/readout_leaf_context.py

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

from __future__ import annotations

import math

import equinox as eqx
import jax
import jax.numpy as jnp
from jaxtyping import Array, Float, Int, PRNGKeyArray

from .odd_ops import BiasFreeLinear, MLP, _RMS


def _inline_norm(
    norm,
    x,
    *,
    pathway="even",
    kfac_structural_mask=None,
    kfac_scan_shared: bool = False,
    kfac_repeat_ndim: int = 0,
    kfac_context_primal_reused_over_walkers: bool = False,
):
    from .tree import _tagged_rms_eqx_style

    return _tagged_rms_eqx_style(
        norm.weight,
        x,
        eps=norm.eps,
        tag_id=norm._use_id,
        pathway=pathway,
        kfac_structural_mask=kfac_structural_mask,
        kfac_scan_shared=kfac_scan_shared,
        kfac_repeat_ndim=kfac_repeat_ndim,
        kfac_context_primal_reused_over_walkers=(
            kfac_context_primal_reused_over_walkers
        ),
    )


def _inline_mlp_forward(
    mlp,
    x,
    *,
    pathway="even",
    kfac_structural_mask=None,
    kfac_scan_shared: bool = False,
    kfac_repeat_ndim: int = 0,
    kfac_context_primal_reused_over_walkers: bool = False,
):
    from .tree import _tagged_dense

    arguments = dict(
        pathway=pathway,
        kfac_structural_mask=kfac_structural_mask,
        kfac_scan_shared=kfac_scan_shared,
        kfac_repeat_ndim=kfac_repeat_ndim,
        kfac_context_primal_reused_over_walkers=(
            kfac_context_primal_reused_over_walkers
        ),
    )
    x = _tagged_dense(
        mlp.in_proj.weight,
        mlp.in_proj.bias,
        x,
        tag_id=mlp.in_proj._use_id,
        **arguments,
    )
    for norm, linear_1, linear_2 in zip(
        mlp.block_norms,
        mlp.block_l1s,
        mlp.block_l2s,
        strict=True,
    ):
        normalized = _inline_norm(norm, x, **arguments)
        inner = _tagged_dense(
            linear_1.weight,
            linear_1.bias,
            normalized,
            tag_id=linear_1._use_id,
            **arguments,
        )
        inner = mlp._act(inner)
        inner = _tagged_dense(
            linear_2.weight,
            linear_2.bias,
            inner,
            tag_id=linear_2._use_id,
            **arguments,
        )
        x = x + mlp.inner_gain * inner
    normalized = _inline_norm(mlp.out_norm, x, **arguments)
    return _tagged_dense(
        mlp.out_proj.weight,
        mlp.out_proj.bias,
        normalized,
        tag_id=mlp.out_proj._use_id,
        **arguments,
    )


def _inline_bias_free_linear(
    layer: BiasFreeLinear,
    x,
    *,
    pathway="even",
    kfac_structural_mask=None,
    kfac_scan_shared: bool = False,
    kfac_repeat_ndim: int = 0,
    kfac_context_primal_reused_over_walkers: bool = False,
):
    from .tree import _tagged_dense_no_bias

    return _tagged_dense_no_bias(
        layer.weight,
        x,
        tag_id=layer._use_id,
        pathway=pathway,
        kfac_structural_mask=kfac_structural_mask,
        kfac_scan_shared=kfac_scan_shared,
        kfac_repeat_ndim=kfac_repeat_ndim,
        kfac_context_primal_reused_over_walkers=(
            kfac_context_primal_reused_over_walkers
        ),
    )


_ALLOWED_ATTN_IMPLS = ("einsum", "mhsea_tuned")
_LCA_MAX_LEVELS = 13


def default_tree_depth(max_n: int) -> int:
    max_n = max(1, int(max_n))
    return max(1, int(math.ceil(math.log2(max_n))))


def _relative_positions(n: int) -> tuple[Int[Array, "n n"], Float[Array, "n n 1"]]:
    index = jnp.arange(n, dtype=jnp.int32)
    relative = index[:, None] - index[None, :]
    sign = jnp.where(
        relative < 0,
        1.0,
        jnp.where(relative > 0, -1.0, 0.0),
    )
    return relative, sign[..., None]


def lca_level(pos_q, pos_k):
    xor = jnp.bitwise_xor(
        jnp.asarray(pos_q, jnp.int32)[..., :, None],
        jnp.asarray(pos_k, jnp.int32)[..., None, :],
    )
    thresholds = jnp.exp2(jnp.arange(_LCA_MAX_LEVELS, dtype=jnp.float32)).astype(
        jnp.int32
    )
    return jnp.sum((xor[..., None] >= thresholds).astype(jnp.int32), axis=-1)


def lca_alibi_bias(pos_q, pos_k, slopes):
    level = lca_level(pos_q, pos_k).astype(slopes.dtype)
    return -(slopes[:, None, None] * level[None, :, :])


def lca_fixed_slopes(n_heads, dtype=jnp.float32):
    head = jnp.arange(n_heads, dtype=jnp.float32)
    denominator = jnp.maximum(jnp.asarray(n_heads - 1, jnp.float32), 1.0)
    return (1.5 * jnp.exp2(-(head * 4.0 / denominator))).astype(dtype)


def lca_gaussian_decay(pos_q, pos_k, w_raw, b):
    level = lca_level(pos_q, pos_k).astype(b.dtype)
    weight = jax.nn.softplus(w_raw)
    distance = level[:, :, None] - b[None, None, :]
    return jnp.exp(-(weight[None, None, :] * distance * distance))


def lca_gaussian_decay_row(pos_q_scalar, pos_k, w_raw, b):
    xor = jnp.bitwise_xor(
        jnp.asarray(pos_q_scalar, jnp.int32),
        jnp.asarray(pos_k, jnp.int32),
    )
    thresholds = jnp.exp2(jnp.arange(_LCA_MAX_LEVELS, dtype=jnp.float32)).astype(
        jnp.int32
    )
    level = jnp.sum(
        (xor[:, None] >= thresholds[None, :]).astype(jnp.int32),
        axis=-1,
    ).astype(b.dtype)
    weight = jax.nn.softplus(w_raw)
    distance = level[:, None] - b[None, :]
    return jnp.exp(-(weight[None, :] * distance * distance))


def register_vector_as_dense(w2d, *, tag_id=""):
    from hamiltonzero.optim.spin_blocks import register_small_full

    return register_small_full(w2d, tag_id=tag_id)


def lca_order_init_w_b(d_model, dtype=jnp.float32):
    width = int(d_model)
    centers = (
        (jnp.arange(width, dtype=jnp.float32) % 9.0).astype(dtype).reshape(1, width)
    )
    raw_weight = jnp.full(
        (1, width),
        float(jnp.log(jnp.expm1(jnp.asarray(0.7)))),
        dtype=dtype,
    )
    return raw_weight, centers


def _attention_dimensions(
    *,
    d_e: int,
    n_heads: int,
    attn_dim: int,
    attn_impl: str,
    require_even_model: bool,
) -> tuple[int, int, int]:
    if n_heads < 1:
        raise ValueError("contextualizer n_heads must be positive")
    if attn_dim < 1 or attn_dim % n_heads:
        raise ValueError("contextualizer attention width must divide by n_heads")
    if require_even_model and d_e % 2:
        raise ValueError("physical contextualizer width must be even")
    d_head = attn_dim // n_heads
    if d_head % 2:
        raise ValueError("contextualizer attention head width must be even")
    if attn_impl not in _ALLOWED_ATTN_IMPLS:
        raise ValueError("attention must be 'einsum' or 'mhsea_tuned'")
    return 2 * n_heads, d_head, attn_dim


def _global_modules(
    key, d_g: int, d_e: int, residual_scale: float, global_tap_dim: int
):
    from .global_ladder import GDescriptorPool, ResidualGlobalUpdate

    keys = jax.random.split(jax.random.fold_in(key, 25007), 3)
    pool = GDescriptorPool(d_g, d_e, key=keys[0], tag="gladder.ctx.pool")
    update = ResidualGlobalUpdate(
        d_g,
        pool.d_out,
        key=keys[1],
        tag="gladder.ctx.upd",
        tap_dim=global_tap_dim,
        residual_gain=residual_scale,
    )
    projection = jax.random.normal(keys[2], (d_g, 64)) * d_g ** (-0.5)
    return pool, update, projection


def _run_attention(query, key, value, bias, mask, *, implementation: str, d_head: int):
    from .pallas_attention import mhsea_tuned_edge_attention, reference_edge_attention

    if implementation == "einsum":
        return reference_edge_attention(query, key, value, bias, mask)
    padded_width = max(16, d_head)
    padding = padded_width - d_head
    scale = jnp.sqrt(jnp.asarray(padded_width / d_head, dtype=jnp.float32)).astype(
        query.dtype
    )
    zeros = jnp.zeros(
        (query.shape[0], query.shape[1], padding),
        dtype=query.dtype,
    )
    query = jnp.concatenate([query * scale, zeros], axis=-1)
    key = jnp.concatenate([key, zeros], axis=-1)
    value = jnp.concatenate([value, zeros], axis=-1)
    return mhsea_tuned_edge_attention(query, key, value, bias, mask)[..., :d_head]


def _edge_update_dense(layer, edge, edge_n, edge_context, bmask_f, g):
    n = edge.shape[0]
    structural = bmask_f.astype(bool)
    pair_structural = structural[:, None] & structural[None, :]
    context_i = jnp.broadcast_to(
        edge_context[:, None, :],
        (n, n, edge_context.shape[-1]),
    )
    context_j = jnp.broadcast_to(
        edge_context[None, :, :],
        (n, n, edge_context.shape[-1]),
    )
    pair_mask = (bmask_f[:, None] * bmask_f[None, :])[..., None]
    from .tree import _tagged_dense_no_bias

    global_edge = _tagged_dense_no_bias(
        layer.g_edge_proj_w,
        g,
        tag_id="gladder.ctx.eproj",
        pathway="even",
        kfac_structural_mask=jnp.any(structural),
        kfac_repeat_ndim=0,
        kfac_context_primal_reused_over_walkers=True,
    )
    edge_input = jnp.concatenate(
        [
            pair_mask * edge_n,
            context_i,
            context_j,
            jnp.broadcast_to(
                global_edge[None, None, :],
                (n, n, global_edge.shape[-1]),
            ).astype(edge_n.dtype),
        ],
        axis=-1,
    )
    delta = _inline_mlp_forward(
        layer.edge_ffn,
        edge_input,
        pathway="even",
        kfac_structural_mask=pair_structural,
        kfac_repeat_ndim=2,
        kfac_context_primal_reused_over_walkers=True,
    )
    return jnp.where(
        pair_structural[..., None],
        edge + layer.residual_scale * delta,
        jnp.zeros_like(edge),
    )


def _edge_update_tiled(
    layer,
    edge_rows,
    edge_n_rows,
    edge_context_rows,
    edge_context_all,
    bmask,
    *,
    row_indices=None,
    g,
    tile_size: int,
):
    rows, n, _ = edge_rows.shape
    if edge_n_rows.shape != edge_rows.shape:
        raise ValueError("normalized edge rows must match edge rows")
    if edge_context_all.shape[0] != n or bmask.shape != (n,):
        raise ValueError("context and mask must have global width N")
    if row_indices is None:
        if rows != n:
            raise ValueError("row_indices is required for row-sharded edges")
        row_indices = jnp.arange(n, dtype=jnp.int32)
    row_indices = jnp.asarray(row_indices, dtype=jnp.int32)
    mask = bmask.astype(edge_rows.dtype)
    mask_rows = mask[row_indices]
    from .tree import _tagged_dense_no_bias

    global_edge = _tagged_dense_no_bias(
        layer.g_edge_proj_w,
        g,
        tag_id="gladder.ctx.eproj",
        pathway="even",
        kfac_structural_mask=jnp.any(bmask.astype(bool)),
        kfac_repeat_ndim=0,
        kfac_context_primal_reused_over_walkers=True,
    )
    tile_width = min(n, int(tile_size))
    if tile_width < 1:
        raise ValueError("tile_size must be positive")
    full_tiles = n // tile_width
    tail_start = full_tiles * tile_width

    def update_tile(start, width, output):
        mask_tile = jax.lax.dynamic_slice_in_dim(mask, start, width, axis=0)
        edge_n_tile = jax.lax.dynamic_slice_in_dim(edge_n_rows, start, width, axis=1)
        context_tile = jax.lax.dynamic_slice_in_dim(
            edge_context_all, start, width, axis=0
        )
        edge_tile = jax.lax.dynamic_slice_in_dim(output, start, width, axis=1)
        pair_structural = mask_rows[:, None].astype(bool) & mask_tile[None, :].astype(
            bool
        )
        context_i = jnp.broadcast_to(
            edge_context_rows[:, None, :],
            (rows, width, edge_context_rows.shape[-1]),
        )
        context_j = jnp.broadcast_to(
            context_tile[None, :, :],
            (rows, width, edge_context_all.shape[-1]),
        )
        edge_input = jnp.concatenate(
            [
                (mask_rows[:, None] * mask_tile[None, :])[..., None] * edge_n_tile,
                context_i,
                context_j,
                jnp.broadcast_to(
                    global_edge,
                    (rows, width, global_edge.shape[-1]),
                ).astype(edge_rows.dtype),
            ],
            axis=-1,
        )
        delta = _inline_mlp_forward(
            layer.edge_ffn,
            edge_input,
            pathway="even",
            kfac_structural_mask=pair_structural,
            kfac_repeat_ndim=2,
            kfac_context_primal_reused_over_walkers=True,
        )
        updated = jnp.where(
            pair_structural[..., None],
            edge_tile + layer.residual_scale * delta,
            jnp.zeros_like(edge_tile),
        )
        return jax.lax.dynamic_update_slice_in_dim(output, updated, start, axis=1)

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


class PhysicalReadoutContextLayer(eqx.Module):
    ln_edge: _RMS
    ln_edge_attn: _RMS
    ln_c: _RMS
    ln_summary: _RMS
    ln_attn: _RMS
    ln_edge_ctx: _RMS
    summary_mlp: MLP
    ctx_mlp: MLP
    edge_node_ctx_proj: BiasFreeLinear
    edge_ffn: MLP
    g_pool: "GDescriptorPool"
    g_update: "ResidualGlobalUpdate"
    g_edge_proj_w: Float[Array, "d_g 64"]
    W_QKV: BiasFreeLinear
    W_O: BiasFreeLinear
    bias_mlp: MLP
    d_e: int = eqx.field(static=True)
    n_heads: int = eqx.field(static=True)
    n_heads_kernel: int = eqx.field(static=True)
    d_head: int = eqx.field(static=True)
    d_attn: int = eqx.field(static=True)
    attn_impl: str = eqx.field(static=True)
    rope_base: float = eqx.field(static=True)
    rope_scaling: float = eqx.field(static=True)
    residual_scale: float = eqx.field(static=True)

    def __init__(
        self,
        *,
        d_e: int,
        d_edge: int,
        n_heads: int,
        summary_hidden: int,
        mlp_hidden: int,
        bias_hidden: int,
        edge_ffn_hidden: int,
        attn_dim: int,
        edge_node_ctx_dim: int,
        attn_impl: str,
        rope_base: float,
        rope_scaling: float,
        gladder_d_g: int,
        global_tap_dim: int,
        residual_scale: float,
        key: PRNGKeyArray,
    ):
        n_heads_kernel, d_head, d_attn = _attention_dimensions(
            d_e=d_e,
            n_heads=n_heads,
            attn_dim=attn_dim,
            attn_impl=attn_impl,
            require_even_model=True,
        )
        if rope_base <= 0.0 or rope_scaling <= 0.0:
            raise ValueError("physical contextualizer rope values must be positive")
        key_sum, key_ctx, key_edge, key_qkv, key_out, key_bias = jax.random.split(
            key, 6
        )
        self.ln_edge = _RMS(d_edge)
        self.ln_edge_attn = _RMS(d_edge)
        self.ln_c = _RMS(d_e)
        self.ln_summary = _RMS(d_e)
        self.ln_attn = _RMS(d_e)
        self.ln_edge_ctx = _RMS(d_e)
        self.summary_mlp = MLP(
            2 * d_edge + 1,
            int(summary_hidden),
            d_e,
            key=key_sum,
            n_blocks=1,
        )
        self.ctx_mlp = MLP(
            2 * d_e,
            int(mlp_hidden),
            d_e,
            key=key_ctx,
            n_blocks=2,
        )
        self.edge_node_ctx_proj = BiasFreeLinear(
            d_e,
            int(edge_node_ctx_dim),
            key=jax.random.fold_in(key, 3694),
        )
        self.edge_ffn = MLP(
            d_edge + 2 * int(edge_node_ctx_dim) + 64,
            int(edge_ffn_hidden),
            d_edge,
            key=key_edge,
            n_blocks=1,
        )
        self.g_pool, self.g_update, self.g_edge_proj_w = _global_modules(
            key,
            int(gladder_d_g),
            d_e,
            residual_scale,
            global_tap_dim,
        )
        self.W_QKV = BiasFreeLinear(d_e, 3 * n_heads_kernel * d_head, key=key_qkv)
        self.W_O = BiasFreeLinear(d_attn, d_e, key=key_out)
        self.bias_mlp = MLP(
            d_edge + 1,
            int(bias_hidden),
            n_heads_kernel,
            key=key_bias,
            n_blocks=1,
        )
        self.d_e = int(d_e)
        self.n_heads = int(n_heads)
        self.n_heads_kernel = int(n_heads_kernel)
        self.d_head = int(d_head)
        self.d_attn = int(d_attn)
        self.attn_impl = str(attn_impl)
        self.rope_base = float(rope_base)
        self.rope_scaling = float(rope_scaling)
        self.residual_scale = float(residual_scale)

    def _slot_clock(self, n: int, dtype, bmask) -> Float[Array, "n d_e"]:
        position = jnp.arange(n, dtype=jnp.float32)
        n_active = jnp.maximum(jnp.sum(jnp.asarray(bmask, dtype=jnp.int32)), 1)
        depth = jnp.ceil(jnp.log2(n_active.astype(jnp.float32))).astype(jnp.int32)
        span = jnp.power(
            jnp.asarray(2.0, dtype=jnp.float32),
            depth.astype(jnp.float32),
        )
        position = position - (span - jnp.asarray(1.0, jnp.float32)) * 0.5
        position = position / jnp.asarray(self.rope_scaling, jnp.float32)
        half = (self.d_e + 1) // 2
        band = jnp.arange(half, dtype=jnp.float32)
        inverse_frequency = jnp.exp(
            -jnp.log(jnp.asarray(self.rope_base, jnp.float32))
            * band
            / jnp.asarray(max(half, 1), jnp.float32)
        )
        angle = position[:, None] * inverse_frequency[None, :]
        embedding = jnp.concatenate([jnp.sin(angle), jnp.cos(angle)], axis=-1)
        return embedding[:, : self.d_e].astype(dtype)

    def _edge_summary(self, edge_n, direction, bmask_f):
        n = edge_n.shape[0]
        dtype = edge_n.dtype
        structural = bmask_f.astype(bool)
        pair_structural = (
            structural[:, None] & structural[None, :] & ~jnp.eye(n, dtype=bool)
        )
        pair = jnp.concatenate(
            [edge_n, jnp.swapaxes(edge_n, 0, 1), direction.astype(dtype)],
            axis=-1,
        )
        message = _inline_mlp_forward(
            self.summary_mlp,
            pair,
            pathway="even",
            kfac_structural_mask=pair_structural,
            kfac_repeat_ndim=2,
            kfac_context_primal_reused_over_walkers=True,
        )
        off_diagonal = 1.0 - jnp.eye(n, dtype=dtype)
        weight = bmask_f[:, None] * bmask_f[None, :] * off_diagonal
        index = jnp.arange(n, dtype=jnp.int32)
        weight = weight * jnp.exp(-0.5 * lca_level(index, index).astype(dtype))
        denominator = jnp.maximum(jnp.sum(weight, axis=1, keepdims=True), 1.0)
        return jnp.sum(message * weight[..., None], axis=1) / denominator

    def edge_summary_tiled(
        self,
        edge_n_rows,
        summary_mask,
        *,
        edge_reverse_rows=None,
        row_indices=None,
        tile_size: int = 128,
    ):
        rows, n, _ = edge_n_rows.shape
        dtype = edge_n_rows.dtype
        if summary_mask.shape != (n,):
            raise ValueError("summary_mask must have global shape [N]")
        if row_indices is None:
            if rows != n:
                raise ValueError("row_indices is required for sharded summaries")
            row_indices = jnp.arange(n, dtype=jnp.int32)
        row_indices = jnp.asarray(row_indices, dtype=jnp.int32)
        if edge_reverse_rows is None:
            if rows != n:
                raise ValueError("reverse rows are required for sharded summaries")
            edge_reverse_rows = jnp.swapaxes(edge_n_rows, 0, 1)
        mask = summary_mask.astype(dtype)
        mask_rows = mask[row_indices]
        numerator = jnp.zeros((rows, self.d_e), dtype=dtype)
        denominator = jnp.zeros((rows, 1), dtype=dtype)
        tile_width = min(n, int(tile_size))
        if tile_width < 1:
            raise ValueError("tile_size must be positive")
        full_tiles = n // tile_width
        tail_start = full_tiles * tile_width

        def accumulate(start, width, carry):
            numerator, denominator = carry
            column = start + jnp.arange(width, dtype=jnp.int32)
            relative = row_indices[:, None] - column[None, :]
            direction = jnp.where(
                relative < 0,
                1.0,
                jnp.where(relative > 0, -1.0, 0.0),
            ).astype(dtype)[..., None]
            off_diagonal = row_indices[:, None] != column[None, :]
            mask_tile = jax.lax.dynamic_slice_in_dim(mask, start, width, axis=0)
            edge_tile = jax.lax.dynamic_slice_in_dim(edge_n_rows, start, width, axis=1)
            reverse_tile = jax.lax.dynamic_slice_in_dim(
                edge_reverse_rows, start, width, axis=1
            )
            structural = (
                mask_rows[:, None].astype(bool)
                & mask_tile[None, :].astype(bool)
                & off_diagonal
            )
            pair = jnp.concatenate([edge_tile, reverse_tile, direction], axis=-1)
            message = _inline_mlp_forward(
                self.summary_mlp,
                pair,
                pathway="even",
                kfac_structural_mask=structural,
                kfac_repeat_ndim=2,
                kfac_context_primal_reused_over_walkers=True,
            )
            weight = (
                mask_rows[:, None] * mask_tile[None, :] * off_diagonal.astype(dtype)
            )
            weight = weight * jnp.exp(
                -0.5 * lca_level(row_indices, column).astype(dtype)
            )
            numerator = numerator + jnp.sum(message * weight[..., None], axis=1)
            denominator = denominator + jnp.sum(weight, axis=1, keepdims=True)
            return numerator, denominator

        numerator, denominator = jax.lax.fori_loop(
            0,
            full_tiles,
            lambda tile_index, carry: accumulate(
                tile_index * tile_width, tile_width, carry
            ),
            (numerator, denominator),
        )
        if tail_start < n:
            numerator, denominator = accumulate(
                tail_start,
                n - tail_start,
                (numerator, denominator),
            )
        return numerator / jnp.maximum(denominator, 1.0)

    def edge_update_tiled(self, *args, **kwargs):
        return _edge_update_tiled(self, *args, **kwargs)

    def _attend(self, c, edge_n, direction, bmask):
        n = c.shape[0]
        dtype = c.dtype
        structural = bmask.astype(bool)
        pair_structural = structural[:, None] & structural[None, :]
        qkv = _inline_bias_free_linear(
            self.W_QKV,
            c,
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        ).reshape(n, 3, self.n_heads_kernel, self.d_head)
        query, key, value = qkv[:, 0], qkv[:, 1], qkv[:, 2]
        bias_input = jnp.concatenate([edge_n, direction.astype(dtype)], axis=-1)
        bias = _inline_mlp_forward(
            self.bias_mlp,
            bias_input,
            pathway="even",
            kfac_structural_mask=pair_structural,
            kfac_repeat_ndim=2,
            kfac_context_primal_reused_over_walkers=True,
        )
        bias = bias / jnp.sqrt(jnp.asarray(self.d_head, dtype=dtype))
        index = jnp.arange(n, dtype=jnp.int32)
        bias = bias + jnp.transpose(
            lca_alibi_bias(
                index,
                index,
                lca_fixed_slopes(self.n_heads_kernel, dtype=dtype),
            ),
            (1, 2, 0),
        )
        output = _run_attention(
            query,
            key,
            value,
            bias,
            bmask.astype(dtype),
            implementation=self.attn_impl,
            d_head=self.d_head,
        )
        output = jax.nn.sigmoid(output[:, : self.n_heads]) * output[:, self.n_heads :]
        return _inline_bias_free_linear(
            self.W_O,
            output.reshape(n, -1),
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        )

    def __call__(self, c, edge, mask, bmask, g):
        del mask
        n = c.shape[0]
        dtype = c.dtype
        bmask_f = bmask.astype(dtype)
        structural = bmask.astype(bool)
        pair_structural = structural[:, None] & structural[None, :]
        edge_n = _inline_norm(
            self.ln_edge,
            edge.astype(dtype),
            pathway="even",
            kfac_structural_mask=pair_structural,
            kfac_repeat_ndim=2,
            kfac_context_primal_reused_over_walkers=True,
        )
        _relative, direction = _relative_positions(n)
        direction = direction.astype(dtype)
        clock = self._slot_clock(n, dtype, bmask)
        summary = self._edge_summary(edge_n, direction, bmask_f)
        context_input = jnp.concatenate(
            [
                _inline_norm(
                    self.ln_c,
                    c + clock,
                    pathway="even",
                    kfac_structural_mask=structural,
                    kfac_repeat_ndim=1,
                    kfac_context_primal_reused_over_walkers=True,
                ),
                _inline_norm(
                    self.ln_summary,
                    summary,
                    pathway="even",
                    kfac_structural_mask=structural,
                    kfac_repeat_ndim=1,
                    kfac_context_primal_reused_over_walkers=True,
                ),
            ],
            axis=-1,
        )
        delta = self.residual_scale * _inline_mlp_forward(
            self.ctx_mlp,
            context_input,
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        )
        c1 = jnp.where(structural[:, None], c + delta, jnp.zeros_like(c))
        global_active = jnp.any(structural)
        g = self.g_update(
            g,
            self.g_pool(
                g,
                c1,
                bmask_f,
                kfac_structural_mask=structural,
                kfac_update_mask=global_active,
                kfac_repeat_ndim=1,
                kfac_context_primal_reused_over_walkers=True,
            ),
            kfac_structural_mask=global_active,
            kfac_context_primal_reused_over_walkers=True,
        )
        edge_context = _inline_bias_free_linear(
            self.edge_node_ctx_proj,
            _inline_norm(
                self.ln_edge_ctx,
                c1,
                pathway="even",
                kfac_structural_mask=structural,
                kfac_repeat_ndim=1,
                kfac_context_primal_reused_over_walkers=True,
            ),
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        )
        edge1 = _edge_update_dense(self, edge, edge_n, edge_context, bmask_f, g)
        attention_input = _inline_norm(
            self.ln_attn,
            c1 + clock,
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        )
        edge_attention = _inline_norm(
            self.ln_edge_attn,
            edge1.astype(dtype),
            pathway="even",
            kfac_structural_mask=pair_structural,
            kfac_repeat_ndim=2,
            kfac_context_primal_reused_over_walkers=True,
        )
        delta_attention = self.residual_scale * self._attend(
            attention_input, edge_attention, direction, bmask
        )
        output = jnp.where(
            structural[:, None],
            c1 + delta_attention,
            jnp.zeros_like(c1),
        )
        return output, edge1, g


class RouterContextLayer(eqx.Module):
    ln_edge: _RMS
    ln_edge_attn: _RMS
    ln_c: _RMS
    ln_attn: _RMS
    ln_edge_ctx: _RMS
    ctx_mlp: MLP
    edge_node_ctx_proj: BiasFreeLinear
    edge_ffn: MLP
    g_pool: "GDescriptorPool"
    g_update: "ResidualGlobalUpdate"
    g_edge_proj_w: Float[Array, "d_g 64"]
    W_QKV: BiasFreeLinear
    W_O: BiasFreeLinear
    bias_mlp: MLP
    n_heads: int = eqx.field(static=True)
    n_heads_kernel: int = eqx.field(static=True)
    d_head: int = eqx.field(static=True)
    d_attn: int = eqx.field(static=True)
    attn_impl: str = eqx.field(static=True)
    residual_scale: float = eqx.field(static=True)

    def __init__(
        self,
        *,
        d_e: int,
        d_edge: int,
        n_heads: int,
        mlp_hidden: int,
        bias_hidden: int,
        edge_ffn_hidden: int,
        attn_dim: int,
        edge_node_ctx_dim: int,
        attn_impl: str,
        gladder_d_g: int,
        global_tap_dim: int,
        residual_scale: float,
        key: PRNGKeyArray,
    ):
        n_heads_kernel, d_head, d_attn = _attention_dimensions(
            d_e=d_e,
            n_heads=n_heads,
            attn_dim=attn_dim,
            attn_impl=attn_impl,
            require_even_model=False,
        )
        _key_sum, key_ctx, key_edge, key_qkv, key_out, key_bias = jax.random.split(
            key, 6
        )
        self.ln_edge = _RMS(d_edge)
        self.ln_edge_attn = _RMS(d_edge)
        self.ln_c = _RMS(d_e)
        self.ln_attn = _RMS(d_e)
        self.ln_edge_ctx = _RMS(d_e)
        self.ctx_mlp = MLP(
            d_e,
            int(mlp_hidden),
            d_e,
            key=key_ctx,
            n_blocks=2,
        )
        self.edge_node_ctx_proj = BiasFreeLinear(
            d_e,
            int(edge_node_ctx_dim),
            key=jax.random.fold_in(key, 3694),
        )
        self.edge_ffn = MLP(
            d_edge + 2 * int(edge_node_ctx_dim) + 64,
            int(edge_ffn_hidden),
            d_edge,
            key=key_edge,
            n_blocks=1,
        )
        self.g_pool, self.g_update, self.g_edge_proj_w = _global_modules(
            key,
            int(gladder_d_g),
            d_e,
            residual_scale,
            global_tap_dim,
        )
        self.W_QKV = BiasFreeLinear(d_e, 3 * n_heads_kernel * d_head, key=key_qkv)
        self.W_O = BiasFreeLinear(d_attn, d_e, key=key_out)
        self.bias_mlp = MLP(
            d_edge + 1,
            int(bias_hidden),
            n_heads_kernel,
            key=key_bias,
            n_blocks=1,
        )
        self.n_heads = int(n_heads)
        self.n_heads_kernel = int(n_heads_kernel)
        self.d_head = int(d_head)
        self.d_attn = int(d_attn)
        self.attn_impl = str(attn_impl)
        self.residual_scale = float(residual_scale)

    def edge_update_tiled(self, *args, **kwargs):
        return _edge_update_tiled(self, *args, **kwargs)

    def _attend(self, c, edge_n, bmask):
        n = c.shape[0]
        dtype = c.dtype
        structural = bmask.astype(bool)
        pair_structural = structural[:, None] & structural[None, :]
        qkv = _inline_bias_free_linear(
            self.W_QKV,
            c,
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        ).reshape(n, 3, self.n_heads_kernel, self.d_head)
        query, key, value = qkv[:, 0], qkv[:, 1], qkv[:, 2]
        direction = jnp.zeros((n, n, 1), dtype=dtype)
        bias = _inline_mlp_forward(
            self.bias_mlp,
            jnp.concatenate([edge_n, direction], axis=-1),
            pathway="even",
            kfac_structural_mask=pair_structural,
            kfac_repeat_ndim=2,
            kfac_context_primal_reused_over_walkers=True,
        )
        bias = bias / jnp.sqrt(jnp.asarray(self.d_head, dtype=dtype))
        output = _run_attention(
            query,
            key,
            value,
            bias,
            bmask.astype(dtype),
            implementation=self.attn_impl,
            d_head=self.d_head,
        )
        output = jax.nn.sigmoid(output[:, : self.n_heads]) * output[:, self.n_heads :]
        return _inline_bias_free_linear(
            self.W_O,
            output.reshape(n, -1),
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        )

    def __call__(self, c, edge, mask, bmask, g):
        del mask
        dtype = c.dtype
        bmask_f = bmask.astype(dtype)
        structural = bmask.astype(bool)
        pair_structural = structural[:, None] & structural[None, :]
        edge_n = _inline_norm(
            self.ln_edge,
            edge.astype(dtype),
            pathway="even",
            kfac_structural_mask=pair_structural,
            kfac_repeat_ndim=2,
            kfac_context_primal_reused_over_walkers=True,
        )
        context_input = _inline_norm(
            self.ln_c,
            c,
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        )
        delta = self.residual_scale * _inline_mlp_forward(
            self.ctx_mlp,
            context_input,
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        )
        c1 = jnp.where(structural[:, None], c + delta, jnp.zeros_like(c))
        global_active = jnp.any(structural)
        g = self.g_update(
            g,
            self.g_pool(
                g,
                c1,
                bmask_f,
                kfac_structural_mask=structural,
                kfac_update_mask=global_active,
                kfac_repeat_ndim=1,
                kfac_context_primal_reused_over_walkers=True,
            ),
            kfac_structural_mask=global_active,
            kfac_context_primal_reused_over_walkers=True,
        )
        edge_context = _inline_bias_free_linear(
            self.edge_node_ctx_proj,
            _inline_norm(
                self.ln_edge_ctx,
                c1,
                pathway="even",
                kfac_structural_mask=structural,
                kfac_repeat_ndim=1,
                kfac_context_primal_reused_over_walkers=True,
            ),
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        )
        edge1 = _edge_update_dense(self, edge, edge_n, edge_context, bmask_f, g)
        attention_input = _inline_norm(
            self.ln_attn,
            c1,
            pathway="even",
            kfac_structural_mask=structural,
            kfac_repeat_ndim=1,
            kfac_context_primal_reused_over_walkers=True,
        )
        edge_attention = _inline_norm(
            self.ln_edge_attn,
            edge1.astype(dtype),
            pathway="even",
            kfac_structural_mask=pair_structural,
            kfac_repeat_ndim=2,
            kfac_context_primal_reused_over_walkers=True,
        )
        delta_attention = self.residual_scale * self._attend(
            attention_input, edge_attention, bmask
        )
        output = jnp.where(
            structural[:, None],
            c1 + delta_attention,
            jnp.zeros_like(c1),
        )
        return output, edge1, g


def _stack_layers(layers):
    dynamic_static = [eqx.partition(layer, eqx.is_array) for layer in layers]
    dynamic = [item[0] for item in dynamic_static]
    static = dynamic_static[0][1]
    stacked = jax.tree.map(lambda *values: jnp.stack(values, axis=0), *dynamic)
    return eqx.combine(stacked, static)


def _initialize_structural_inputs(contextualizer, c, edge, mask, bmask):
    real = mask.astype(bool)
    active = bmask.astype(bool)
    virtual = active & ~real
    virtual_node = register_vector_as_dense(
        contextualizer.virtual_node,
        tag_id=contextualizer._use_id_virtual_node,
    )[0].astype(c.dtype)
    empty_nonempty = register_vector_as_dense(
        contextualizer.edge_empty_nonempty,
        tag_id=contextualizer._use_id_edge_empty_nonempty,
    )[0].astype(edge.dtype)
    empty_empty = register_vector_as_dense(
        contextualizer.edge_empty_empty,
        tag_id=contextualizer._use_id_edge_empty_empty,
    )[0].astype(edge.dtype)
    c = jnp.where(
        real[:, None],
        c,
        jnp.where(
            virtual[:, None],
            virtual_node[None, :],
            jnp.zeros_like(c),
        ),
    )
    real_pair = real[:, None] & real[None, :]
    mixed_pair = (real[:, None] & virtual[None, :]) | (virtual[:, None] & real[None, :])
    virtual_pair = virtual[:, None] & virtual[None, :]
    edge = jnp.where(
        real_pair[..., None],
        edge,
        jnp.where(
            mixed_pair[..., None],
            empty_nonempty[None, None, :],
            jnp.where(
                virtual_pair[..., None],
                empty_empty[None, None, :],
                jnp.zeros_like(edge),
            ),
        ),
    )
    return c, edge


def _run_context_layers(contextualizer, c, edge, mask, bmask, g):
    dynamic, static = eqx.partition(contextualizer.layers, eqx.is_array)

    def scan_step(carry, layer_dynamic):
        layer = eqx.combine(layer_dynamic, static)
        c_value, edge_value, g_value = carry
        return (
            layer(c_value, edge_value, mask, bmask, g_value),
            None,
        )

    (c, edge, g), _ = jax.lax.scan(
        scan_step,
        (c, edge, g),
        dynamic,
    )
    return c, edge, g


class PhysicalReadoutContext(eqx.Module):
    layers: PhysicalReadoutContextLayer
    virtual_node: Float[Array, "one d_e"]
    edge_empty_nonempty: Float[Array, "one d_edge"]
    edge_empty_empty: Float[Array, "one d_edge"]
    _use_id_virtual_node: str = eqx.field(static=True, default="")
    _use_id_edge_empty_nonempty: str = eqx.field(static=True, default="")
    _use_id_edge_empty_empty: str = eqx.field(static=True, default="")

    def __init__(
        self,
        *,
        d_e: int,
        d_edge: int,
        n_layers: int,
        n_heads: int,
        summary_hidden: int,
        mlp_hidden: int,
        bias_hidden: int,
        edge_ffn_hidden: int,
        attn_dim: int,
        edge_node_ctx_dim: int,
        attn_impl: str,
        rope_base: float,
        rope_scaling: float,
        gladder_d_g: int,
        global_tap_dim: int,
        key: PRNGKeyArray,
    ):
        if n_layers < 1:
            raise ValueError("physical contextualizer layers must be positive")
        residual_scale = float(n_layers) ** (-0.5)
        keys = jax.random.split(key, n_layers)
        self.layers = _stack_layers(
            [
                PhysicalReadoutContextLayer(
                    d_e=d_e,
                    d_edge=d_edge,
                    n_heads=n_heads,
                    summary_hidden=summary_hidden,
                    mlp_hidden=mlp_hidden,
                    bias_hidden=bias_hidden,
                    edge_ffn_hidden=edge_ffn_hidden,
                    attn_dim=attn_dim,
                    edge_node_ctx_dim=edge_node_ctx_dim,
                    attn_impl=attn_impl,
                    rope_base=rope_base,
                    rope_scaling=rope_scaling,
                    gladder_d_g=gladder_d_g,
                    global_tap_dim=global_tap_dim,
                    residual_scale=residual_scale,
                    key=layer_key,
                )
                for layer_key in keys
            ]
        )
        self.virtual_node = jax.random.normal(
            jax.random.fold_in(key, 201793223), (1, d_e)
        ) * d_e ** (-0.5)
        self.edge_empty_nonempty = jax.random.normal(
            jax.random.fold_in(key, 235798529), (1, d_edge)
        ) * d_edge ** (-0.5)
        self.edge_empty_empty = jax.random.normal(
            jax.random.fold_in(key, 235798530), (1, d_edge)
        ) * d_edge ** (-0.5)
        self._use_id_virtual_node = ""
        self._use_id_edge_empty_nonempty = ""
        self._use_id_edge_empty_empty = ""

    def with_edge(self, c, edge, mask, bmask, *, g):
        c = c.astype(jnp.float32)
        edge = edge.astype(jnp.float32)
        c, edge = _initialize_structural_inputs(self, c, edge, mask, bmask)
        return _run_context_layers(self, c, edge, mask, bmask, g)


class RouterContext(eqx.Module):
    layers: RouterContextLayer
    virtual_node: Float[Array, "one d_e"]
    edge_empty_nonempty: Float[Array, "one d_edge"]
    edge_empty_empty: Float[Array, "one d_edge"]
    _use_id_virtual_node: str = eqx.field(static=True, default="")
    _use_id_edge_empty_nonempty: str = eqx.field(static=True, default="")
    _use_id_edge_empty_empty: str = eqx.field(static=True, default="")

    def __init__(
        self,
        *,
        d_e: int,
        d_edge: int,
        n_layers: int,
        n_heads: int,
        mlp_hidden: int,
        bias_hidden: int,
        edge_ffn_hidden: int,
        attn_dim: int,
        edge_node_ctx_dim: int,
        attn_impl: str,
        gladder_d_g: int,
        global_tap_dim: int,
        key: PRNGKeyArray,
    ):
        if n_layers < 1:
            raise ValueError("router contextualizer layers must be positive")
        residual_scale = float(n_layers) ** (-0.5)
        keys = jax.random.split(key, n_layers)
        self.layers = _stack_layers(
            [
                RouterContextLayer(
                    d_e=d_e,
                    d_edge=d_edge,
                    n_heads=n_heads,
                    mlp_hidden=mlp_hidden,
                    bias_hidden=bias_hidden,
                    edge_ffn_hidden=edge_ffn_hidden,
                    attn_dim=attn_dim,
                    edge_node_ctx_dim=edge_node_ctx_dim,
                    attn_impl=attn_impl,
                    gladder_d_g=gladder_d_g,
                    global_tap_dim=global_tap_dim,
                    residual_scale=residual_scale,
                    key=layer_key,
                )
                for layer_key in keys
            ]
        )
        self.virtual_node = jax.random.normal(
            jax.random.fold_in(key, 201793223), (1, d_e)
        ) * d_e ** (-0.5)
        self.edge_empty_nonempty = jax.random.normal(
            jax.random.fold_in(key, 235798529), (1, d_edge)
        ) * d_edge ** (-0.5)
        self.edge_empty_empty = jax.random.normal(
            jax.random.fold_in(key, 235798530), (1, d_edge)
        ) * d_edge ** (-0.5)
        self._use_id_virtual_node = ""
        self._use_id_edge_empty_nonempty = ""
        self._use_id_edge_empty_empty = ""

    def with_edge(self, c, edge, mask, bmask, *, g):
        c = c.astype(jnp.float32)
        edge = edge.astype(jnp.float32)
        c, edge = _initialize_structural_inputs(self, c, edge, mask, bmask)
        return _run_context_layers(self, c, edge, mask, bmask, g)


__all__ = [
    "PhysicalReadoutContext",
    "PhysicalReadoutContextLayer",
    "RouterContext",
    "RouterContextLayer",
    "default_tree_depth",
    "lca_alibi_bias",
    "lca_fixed_slopes",
    "lca_gaussian_decay",
    "lca_gaussian_decay_row",
    "lca_level",
    "lca_order_init_w_b",
    "register_vector_as_dense",
]