# Copyright (c) 2026 Simulacra Research Inc. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations from functools import partial import equinox as eqx import jax import jax.numpy as jnp import kfac_jax from jaxtyping import Array, Float, Int, PRNGKeyArray def _kfac_name_kw(tag_id: str) -> dict: return {"name": tag_id} if tag_id else {} _TREE_DYADIC_CLOCK_BASE = 10000.0 def _tree_coord_clock( pos, width: int, dtype, *, base: float = _TREE_DYADIC_CLOCK_BASE, scale=None ): if int(width) <= 0: pos_arr = jnp.asarray(pos) return jnp.zeros(pos_arr.shape + (0,), dtype=dtype) pos_f = jnp.asarray(pos, dtype=jnp.float32) if scale is not None: denom = jnp.maximum( jnp.asarray(scale, dtype=jnp.float32) - jnp.asarray(1.0, dtype=jnp.float32), jnp.asarray(1.0, dtype=jnp.float32), ) pos_f = pos_f / denom half = (int(width) + 1) // 2 band = jnp.arange(half, dtype=jnp.float32) base_f = jnp.maximum( jnp.asarray(base, dtype=jnp.float32), jnp.asarray(2.0, dtype=jnp.float32) ) inv_freq = jnp.exp( -jnp.log(base_f) * band / jnp.asarray(max(half, 1), dtype=jnp.float32) ) phase = pos_f[..., None] * inv_freq enc = jnp.concatenate([jnp.sin(phase), jnp.cos(phase)], axis=-1) return enc[..., : int(width)].astype(dtype) def _tree_level_clock(pos, width: int, max_depth, dtype): del max_depth return _tree_coord_clock(pos, width, dtype, base=_TREE_DYADIC_CLOCK_BASE) def _tree_clock_root_center_from_depth(depth, dtype): depth_f = jnp.maximum(jnp.asarray(depth, dtype=jnp.float32), 0.0) span = jnp.power(jnp.asarray(2.0, dtype=jnp.float32), depth_f) return (span - jnp.asarray(1.0, dtype=jnp.float32)) * 0.5 def _tree_dyadic_segment_clock( level_idx, pair_idx, width: int, dtype, *, root_center=None ): if int(width) <= 0: level_arr = jnp.asarray(level_idx) if pair_idx is None: return jnp.zeros(level_arr.shape + (0,), dtype=dtype) pair_arr = jnp.asarray(pair_idx) return jnp.zeros(pair_arr.shape + (0,), dtype=dtype) center_width = max(1, int(width) - 2) scale_width = int(width) - center_width level_f = jnp.asarray(level_idx, dtype=jnp.float32) span = jnp.power( jnp.asarray(2.0, dtype=jnp.float32), level_f + jnp.asarray(1.0, dtype=jnp.float32), ) scale_pos = span if pair_idx is None: center_pos = (span - jnp.asarray(1.0, dtype=jnp.float32)) * 0.5 else: pair_f = jnp.asarray(pair_idx, dtype=jnp.float32) center_pos = pair_f * span + jnp.asarray(0.5, dtype=jnp.float32) * ( span - jnp.asarray(1.0, dtype=jnp.float32) ) if root_center is not None: center_pos = center_pos - jnp.asarray(root_center, dtype=jnp.float32) scale_pos = jnp.broadcast_to(scale_pos, center_pos.shape) center_clock = _tree_coord_clock( center_pos, center_width, dtype, base=_TREE_DYADIC_CLOCK_BASE ) scale_clock = _tree_coord_clock( scale_pos, scale_width, dtype, base=_TREE_DYADIC_CLOCK_BASE ) return jnp.concatenate([center_clock, scale_clock], axis=-1) def _tree_merge_clock(level_idx, pair_idx, pair_base, width: int, max_depth, dtype): del pair_base root_center = _tree_clock_root_center_from_depth(max_depth, dtype) return _tree_dyadic_segment_clock( level_idx, pair_idx, width, dtype, root_center=root_center ) _TREE_NGPT_DEPTH_FEAT_DIM = 32 def _tree_sphere(x, axis=-1): ms = jnp.mean(jnp.square(x), axis=axis, keepdims=True) return x * jax.lax.rsqrt(jnp.maximum(ms, 0.0001)) def _tree_depth_count_features(cnt_a, cnt_b, n_total, level, n_levels, dtype): a = cnt_a.astype(jnp.float32) b = cnt_b.astype(jnp.float32) nt = jnp.asarray(n_total, jnp.float32) la = jnp.log2(1.0 + a) lb = jnp.log2(1.0 + b) rem = jnp.log2(1.0 + jnp.maximum(nt - a - b, 0.0)) counts = jnp.stack([la, lb, rem], axis=-1) omg_c = jnp.pi / 2.0 ** jnp.arange(4, dtype=jnp.float32) ang_c = counts[..., :, None] * omg_c f_c = jnp.concatenate([jnp.sin(ang_c), jnp.cos(ang_c)], axis=-1) f_c = f_c.reshape(f_c.shape[:-2] + (24,)) lv = jnp.asarray(level, jnp.float32) lv_rem = jnp.asarray(n_levels, jnp.float32) - 1.0 - lv levels = jnp.stack( [jnp.broadcast_to(lv, a.shape), jnp.broadcast_to(lv_rem, a.shape)], axis=-1 ) omg_l = jnp.pi / 2.0 ** jnp.arange(2, dtype=jnp.float32) ang_l = levels[..., :, None] * omg_l f_l = jnp.concatenate([jnp.sin(ang_l), jnp.cos(ang_l)], axis=-1) f_l = f_l.reshape(f_l.shape[:-2] + (8,)) return jnp.concatenate([f_c, f_l], axis=-1).astype(dtype) def _tree_ngpt_level_counts(m0, n_pairs, n_levels, dtype, *, feature_n_levels=None): cnt = m0.astype(jnp.float32) n_total = jnp.sum(cnt) if feature_n_levels is None: feature_n_levels = _tree_active_clock_depth(m0) feats = [] for lv in range(n_levels): pairs = cnt.reshape(n_pairs, 2) feats.append( _tree_depth_count_features( pairs[:, 0], pairs[:, 1], n_total, lv, feature_n_levels, dtype ) ) parents = pairs.sum(axis=1) cnt = jnp.concatenate([parents, jnp.zeros_like(parents)], axis=0) return jnp.stack(feats, axis=0) from .odd_ops import BiasFreeLinear, HypernetMatrix, Linear, MLP, _RMS from .readout_leaf_context import lca_alibi_bias, lca_fixed_slopes from .fused_silu import fused_silu def _replace_square_row_column(matrix, index, row_value, column_value): idx = jnp.arange(matrix.shape[0], dtype=jnp.int32) select = idx == jnp.asarray(index, dtype=jnp.int32) diagonal = row_value[index] column_value = jnp.where(select[:, None], diagonal, column_value) matrix = jnp.where(select[:, None, None], row_value[None, :, :], matrix) return jnp.where(select[None, :, None], column_value[:, None, :], matrix) def _quadrilinear_merge( T, u_a, u_b, *, tag_id: str = "", pathway: str | None = None, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 0, ): from ._custom_lap_primitives import custom_lap_active, quadrilinear_merge_p if custom_lap_active(): return quadrilinear_merge_p.bind(T, u_a, u_b) G, d_r = (T.shape[0], T.shape[1]) _odt = jnp.float32 T_param = T T = T if T.dtype == _odt else T.astype(_odt) u_a = u_a if u_a.dtype == _odt else u_a.astype(_odt) u_b = u_b if u_b.dtype == _odt else u_b.astype(_odt) leading = u_a.shape[:-1] u_a_2d = u_a.reshape(*leading, G, d_r) u_b_2d = u_b.reshape(*leading, G, d_r) Tu_a = jnp.einsum("ijkl,...ik->...ijl", T, u_a_2d) y_2d = jnp.einsum("...ijl,...il->...ij", Tu_a, u_b_2d) y = y_2d.reshape(*leading, G * d_r) if kfac_structural_mask is None: return y from hamiltonzero.optim.spin_blocks import register_structural_quadrilinear_merge return register_structural_quadrilinear_merge( y, u_a, u_b, T_param, kfac_structural_mask, scan_shared=kfac_scan_shared, repeat_ndim=kfac_repeat_ndim, **_kfac_name_kw(tag_id), ) def _rownorm_cols(weight): nsq = jnp.sum(jnp.square(weight), axis=0, keepdims=True) return weight * jax.lax.rsqrt(jnp.maximum(nsq, 0.0001)) def _tagged_dense( weight, bias, x, *, tag_id: str = "", pathway: str, weight_eff=None, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 0, kfac_context_primal_reused_over_walkers: bool = False, ): from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype cdtype = _compute_dtype() if pathway in ("even", "hypernet_eside") else jnp.float32 w_src = weight if weight_eff is None else weight_eff w_compute = w_src.astype(cdtype) if w_src.dtype != cdtype else w_src b_compute = bias.astype(cdtype) if bias.dtype != cdtype else bias x_compute = x.astype(cdtype) if x.dtype != cdtype else x y = x_compute @ w_compute + b_compute if kfac_structural_mask is not None: from hamiltonzero.optim.blocks import register_structural_dense return register_structural_dense( y, x, kfac_structural_mask, weight, bias, scan_shared=kfac_scan_shared, repeat_ndim=kfac_repeat_ndim, context_primal_reused_over_walkers=kfac_context_primal_reused_over_walkers, **_kfac_name_kw(tag_id), ) return kfac_jax.register_dense(y, x, weight, bias, **_kfac_name_kw(tag_id)) def _tagged_dense_no_bias( weight, x, *, tag_id: str = "", pathway: str, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 0, kfac_context_primal_reused_over_walkers: bool = False, ): from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype cdtype = _compute_dtype() if pathway in ("even", "hypernet_eside") else jnp.float32 w_compute = weight.astype(cdtype) if weight.dtype != cdtype else weight x_compute = x.astype(cdtype) if x.dtype != cdtype else x y = x_compute @ w_compute if kfac_structural_mask is not None: from hamiltonzero.optim.blocks import register_structural_dense return register_structural_dense( y, x, kfac_structural_mask, weight, scan_shared=kfac_scan_shared, repeat_ndim=kfac_repeat_ndim, context_primal_reused_over_walkers=kfac_context_primal_reused_over_walkers, **_kfac_name_kw(tag_id), ) return kfac_jax.register_dense(y, x, weight, **_kfac_name_kw(tag_id)) def _tagged_ln_eqx_style( scale, shift, x, eps: float = 1e-05, *, tag_id: str = "", pathway: str, var_floor: float | None = None, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 0, kfac_context_primal_reused_over_walkers: bool = False, ): from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype out_cdtype = ( _compute_dtype() if pathway in ("even", "hypernet_eside") else jnp.float32 ) in_dtype = x.dtype stats_dtype = jnp.promote_types(jnp.float32, in_dtype) x_hi = x.astype(stats_dtype) if in_dtype != stats_dtype else x mean = jnp.mean(x_hi, axis=-1, keepdims=True) centered_hi = x_hi - mean var = jnp.mean(centered_hi * centered_hi, axis=-1, keepdims=True) if var_floor is not None: var = jnp.maximum(var, var_floor) normalized_hi = centered_hi * jax.lax.rsqrt(var + eps) normalized = ( normalized_hi.astype(out_cdtype) if normalized_hi.dtype != out_cdtype else normalized_hi ) scale_compute = scale.astype(out_cdtype) if scale.dtype != out_cdtype else scale shift_compute = shift.astype(out_cdtype) if shift.dtype != out_cdtype else shift y = normalized * scale_compute + shift_compute if kfac_structural_mask is not None: from hamiltonzero.optim.blocks import register_structural_scale_and_shift return register_structural_scale_and_shift( y, normalized, kfac_structural_mask, scale=scale, shift=shift, scan_shared=kfac_scan_shared, repeat_ndim=kfac_repeat_ndim, context_primal_reused_over_walkers=kfac_context_primal_reused_over_walkers, **_kfac_name_kw(tag_id), ) return kfac_jax.register_scale_and_shift( y, normalized, scale, shift, **_kfac_name_kw(tag_id) ) def _tagged_rms_eqx_style( scale, x, eps: float = 1e-05, *, tag_id: str = "", pathway: str, var_floor: float | None = None, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 0, kfac_context_primal_reused_over_walkers: bool = False, ): from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype out_cdtype = ( _compute_dtype() if pathway in ("even", "hypernet_eside") else jnp.float32 ) in_dtype = x.dtype _stats_dtype = jnp.promote_types(jnp.float32, in_dtype) x_hi = x.astype(_stats_dtype) if in_dtype != _stats_dtype else x mean_sq = jnp.mean(x_hi * x_hi, axis=-1, keepdims=True) if var_floor is not None: rsqrt_hi = jax.lax.rsqrt(jnp.maximum(mean_sq, var_floor)) else: rsqrt_hi = jax.lax.rsqrt(mean_sq + eps) normalized_hi = x_hi * rsqrt_hi normalized = ( normalized_hi.astype(out_cdtype) if normalized_hi.dtype != out_cdtype else normalized_hi ) rsqrt = rsqrt_hi.astype(out_cdtype) if rsqrt_hi.dtype != out_cdtype else rsqrt_hi s_compute = scale.astype(out_cdtype) if scale.dtype != out_cdtype else scale x_compute = x.astype(out_cdtype) if x.dtype != out_cdtype else x inv = s_compute * rsqrt y = x_compute * inv if kfac_structural_mask is not None: from hamiltonzero.optim.blocks import register_structural_scale_and_shift return register_structural_scale_and_shift( y, normalized, kfac_structural_mask, scale=scale, scan_shared=kfac_scan_shared, repeat_ndim=kfac_repeat_ndim, context_primal_reused_over_walkers=kfac_context_primal_reused_over_walkers, **_kfac_name_kw(tag_id), ) return kfac_jax.register_scale_and_shift( y, normalized, scale=s_compute, shift=None, **_kfac_name_kw(tag_id) ) def _tagged_lerp_alpha( alpha, d, *, tag_id: str = "", pathway: str, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 0, kfac_context_primal_reused_over_walkers: bool = False, ): from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype out_cdtype = ( _compute_dtype() if pathway in ("even", "hypernet_eside") else jnp.float32 ) d_compute = d.astype(out_cdtype) if d.dtype != out_cdtype else d a_compute = alpha.astype(out_cdtype) if alpha.dtype != out_cdtype else alpha y = d_compute * a_compute if kfac_structural_mask is not None: from hamiltonzero.optim.blocks import register_structural_scale_and_shift return register_structural_scale_and_shift( y, d_compute, kfac_structural_mask, scale=alpha, scan_shared=kfac_scan_shared, repeat_ndim=kfac_repeat_ndim, context_primal_reused_over_walkers=kfac_context_primal_reused_over_walkers, **_kfac_name_kw(tag_id), ) return kfac_jax.register_scale_and_shift( y, d_compute, scale=alpha, shift=None, **_kfac_name_kw(tag_id) ) def _tagged_bounded_ngpt_gain( alpha, like, *, max_gain: float = 0.5, tag_id: str = "", pathway: str, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 0, kfac_context_primal_reused_over_walkers: bool = False, ): ones = jax.lax.stop_gradient(like) * jnp.asarray( 0.0, dtype=like.dtype ) + jnp.asarray(1.0, dtype=like.dtype) tagged_alpha = _tagged_lerp_alpha( alpha, ones, tag_id=tag_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, ) return jnp.asarray(max_gain, dtype=tagged_alpha.dtype) * jax.nn.sigmoid( tagged_alpha ) def _tree_ngpt_residual( skip, proposal, alpha, *, max_gain: float, tag_id: str, pathway: str = "even", update_mask=None, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 0, kfac_context_primal_reused_over_walkers: bool = False, ): skip_n = _tree_sphere(skip) proposal_n = _tree_sphere(proposal) direction = proposal_n - skip_n gain = _tagged_bounded_ngpt_gain( alpha, direction, max_gain=max_gain, tag_id=tag_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, ) updated = _tree_sphere(skip_n + gain * direction) if update_mask is None: return updated active = update_mask.astype(bool) while active.ndim < updated.ndim: active = active[..., None] return jnp.where(active, updated, skip) def _inline_norm_forward( nrm, x, *, pathway: str, tag_id=None, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 0, kfac_context_primal_reused_over_walkers: bool = False, ): tid = nrm._use_id if tag_id is None else tag_id return _tagged_rms_eqx_style( nrm.weight, x, eps=nrm.eps, tag_id=tid, 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, ) class LeafBuilder(eqx.Module): P_c: Linear P_u: HypernetMatrix def __init__( self, d_e: int, d_o: int, d_c: int, d_r: int, rank: int, *, key: PRNGKeyArray, d_g: int, leaf_hypernet_rank: int | None = None, d_m_merge: int | None = None, ): keys = jax.random.split(key, 5) ctx_dim = d_e + d_g p_u_rank = leaf_hypernet_rank if leaf_hypernet_rank is not None else rank d_m_eff = d_m_merge if d_m_merge is not None else d_r if d_m_eff % d_r != 0: raise ValueError( f"LeafBuilder: d_m_merge={d_m_merge} must be divisible by d_r={d_r} (HT carrier reshape requires G = d_m_eff // d_r)." ) self.P_c = Linear(d_e, d_c, key=keys[0]) self.P_u = HypernetMatrix(d_o, d_m_eff, ctx_dim, p_u_rank, key=keys[1]) def conditioner_context( self, e: Float[Array, "n d_e"], g_emb: Float[Array, "d_g"] ) -> Float[Array, "n d_ctx"]: n = e.shape[0] g_emb_b = jnp.broadcast_to(g_emb[None, :], (n, g_emb.shape[0])) return jnp.concatenate([e, g_emb_b], axis=-1) def __call__(self, e, z, *, g_emb, kfac_structural_mask, kfac_odd_structural_mask): n = e.shape[0] ctx = self.conditioner_context(e, g_emb) c = _tagged_dense( self.P_c.weight, self.P_c.bias, e, tag_id=self.P_c._use_id, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=False, kfac_repeat_ndim=1, ) u = self.P_u.apply( ctx, z, e_pathway="even", kfac_structural_mask=kfac_odd_structural_mask, kfac_scan_shared=False, kfac_repeat_ndim=1, kfac_context_primal_reused_over_walkers=True, ) return (_tree_sphere(c), u, jnp.zeros((n,), dtype=jnp.float32)) class EdgeMergeOp(eqx.Module): mlp: MLP node_ctx_proj: BiasFreeLinear | None alpha: Float[Array, "d_edge"] ngpt_alpha_max: float = eqx.field(static=True, default=0.5) def __init__( self, d_edge: int, d_c: int, *, key: PRNGKeyArray, alpha_init: float, alpha_max: float, d_hidden: int | None = None, n_blocks: int = 2, edge_node_ctx_dim: int | None = None, ): node_ctx_dim = int(d_c) if edge_node_ctx_dim is None else int(edge_node_ctx_dim) if node_ctx_dim < 1: raise ValueError( f"tree edge_node_ctx_dim must be positive or None, got {edge_node_ctx_dim}" ) self.node_ctx_proj = ( None if node_ctx_dim == int(d_c) else BiasFreeLinear(d_c, node_ctx_dim, key=jax.random.fold_in(key, 60782)) ) d_in = 4 * d_edge + 4 * node_ctx_dim d_hidden_eff = d_hidden if d_hidden is not None else max(d_edge * 2, 64) self.mlp = MLP(d_in, d_hidden_eff, d_edge, key=key, n_blocks=n_blocks) self.ngpt_alpha_max = float(alpha_max) self.alpha = float(alpha_init) * jnp.ones((int(d_edge),)) def apply_skip( self, skip, proposal, *, kfac_structural_mask=None, kfac_scan_shared: bool = False, ): return _tree_ngpt_residual( skip, proposal, self.alpha, max_gain=self.ngpt_alpha_max, tag_id="", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) def __call__( self, e_2i_2j: Float[Array, "d_edge"], e_2i_2j1: Float[Array, "d_edge"], e_2i1_2j: Float[Array, "d_edge"], e_2i1_2j1: Float[Array, "d_edge"], c_2i: Float[Array, "d_c"], c_2i1: Float[Array, "d_c"], c_2j: Float[Array, "d_c"], c_2j1: Float[Array, "d_c"], *, kfac_structural_mask=None, kfac_scan_shared: bool = False, ) -> Float[Array, "d_edge"]: child_ctx = jnp.stack([c_2i, c_2i1, c_2j, c_2j1], axis=0) if self.node_ctx_proj is not None: child_structural_mask = ( None if kfac_structural_mask is None else jnp.broadcast_to(kfac_structural_mask, (4,)) ) child_ctx = _tagged_dense_no_bias( self.node_ctx_proj.weight, child_ctx, tag_id=self.node_ctx_proj._use_id, pathway="even", kfac_structural_mask=child_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=3, ) c_2i, c_2i1, c_2j, c_2j1 = child_ctx mlp_in = jnp.concatenate( [e_2i_2j, e_2i_2j1, e_2i1_2j, e_2i1_2j1, c_2i, c_2i1, c_2j, c_2j1] ) mlp = self.mlp x = _tagged_dense( mlp.in_proj.weight, mlp.in_proj.bias, mlp_in, tag_id=mlp.in_proj._use_id, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) for nrm, l1, l2 in zip(mlp.block_norms, mlp.block_l1s, mlp.block_l2s): normed = _inline_norm_forward( nrm, x, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) inner = _tagged_dense( l1.weight, l1.bias, normed, tag_id=l1._use_id, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) inner_act = mlp._act(inner) inner_out = _tagged_dense( l2.weight, l2.bias, inner_act, tag_id=l2._use_id, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) x = x + mlp.inner_gain * inner_out out_normed = _inline_norm_forward( mlp.out_norm, x, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) return _tagged_dense( mlp.out_proj.weight, mlp.out_proj.bias, out_normed, tag_id=mlp.out_proj._use_id, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) class EdgeFWLUpdate(eqx.Module): ln_edge: _RMS ln_c: _RMS node_ctx_proj: BiasFreeLinear | None psi_L_in: Linear psi_L_out: Linear psi_R_in: Linear psi_R_out: Linear ln_path: _RMS ffn_in: Linear ffn_out: Linear alpha: Float[Array, "d_edge"] ngpt_alpha_max: float = eqx.field(static=True, default=0.5) def __init__( self, d_c: int, d_edge: int, *, key: PRNGKeyArray, alpha_init: float, alpha_max: float, channels: int = 64, edge_node_ctx_dim: int | None = None, ): node_ctx_dim = int(d_c) if edge_node_ctx_dim is None else int(edge_node_ctx_dim) if node_ctx_dim < 1: raise ValueError( f"tree edge_node_ctx_dim must be positive or None, got {edge_node_ctx_dim}" ) d_pair = d_edge + 2 * node_ctx_dim d_psi_hidden = 2 * channels d_ffn_hidden = max(d_edge, 2 * channels) ( k_psi_L_in, k_psi_L_out, k_psi_L_gate, k_psi_R_in, k_psi_R_out, k_psi_R_gate, k_ffn_in, k_ffn_out, ) = jax.random.split(key, 8) self.ln_edge = _RMS(d_edge) self.ln_c = _RMS(d_c) self.node_ctx_proj = ( None if node_ctx_dim == int(d_c) else BiasFreeLinear(d_c, node_ctx_dim, key=jax.random.fold_in(key, 63262)) ) self.psi_L_in = Linear(d_pair, d_psi_hidden, key=k_psi_L_in) self.psi_L_out = Linear(d_psi_hidden, channels, key=k_psi_L_out) self.psi_R_in = Linear(d_pair, d_psi_hidden, key=k_psi_R_in) self.psi_R_out = Linear(d_psi_hidden, channels, key=k_psi_R_out) self.ln_path = _RMS(channels) self.ffn_in = Linear(d_pair + channels, d_ffn_hidden, key=k_ffn_in) self.ffn_out = Linear(d_ffn_hidden, d_edge, key=k_ffn_out) self.ngpt_alpha_max = float(alpha_max) self.alpha = float(alpha_init) * jnp.ones((int(d_edge),)) def _psi_apply( self, pair_ij: Float[Array, "n n d_pair"], which: str, *, kfac_structural_mask=None, kfac_scan_shared: bool = False, kfac_repeat_ndim: int = 2, ) -> Float[Array, "n n C"]: if which == "L": l_in, l_out = (self.psi_L_in, self.psi_L_out) else: l_in, l_out = (self.psi_R_in, self.psi_R_out) hidden = fused_silu( _tagged_dense( l_in.weight, l_in.bias, pair_ij, tag_id=l_in._use_id, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=kfac_repeat_ndim, ) ) return _tagged_dense( l_out.weight, l_out.bias, hidden, tag_id=l_out._use_id, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=kfac_repeat_ndim, ) def __call__( self, edge: Float[Array, "n n d_edge"], c_level: Float[Array, "n d_c"], mask: Float[Array, "n"] | None = None, *, kfac_scan_shared: bool = False, ) -> Float[Array, "n n d_edge"]: n = c_level.shape[0] d_c_dim = c_level.shape[-1] node_structural_mask = mask full_pair_structural_mask = ( None if mask is None else mask[:, None] * mask[None, :] ) edge_ln = _inline_norm_forward( self.ln_edge, edge, pathway="even", kfac_structural_mask=full_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) c_ln = _inline_norm_forward( self.ln_c, c_level, pathway="even", kfac_structural_mask=node_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) c_ctx = ( c_ln if self.node_ctx_proj is None else _tagged_dense_no_bias( self.node_ctx_proj.weight, c_ln, tag_id=self.node_ctx_proj._use_id, pathway="even", kfac_structural_mask=node_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) ) d_c_dim = c_ctx.shape[-1] c_i_b = jnp.broadcast_to(c_ctx[:, None, :], (n, n, d_c_dim)) c_j_b = jnp.broadcast_to(c_ctx[None, :, :], (n, n, d_c_dim)) pair_ij = jnp.concatenate([edge_ln, c_i_b, c_j_b], axis=-1) A = self._psi_apply( pair_ij, "L", kfac_structural_mask=full_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, ) B = self._psi_apply( pair_ij, "R", kfac_structural_mask=full_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, ) if mask is not None: m = mask.astype(A.dtype) A = A * (m[:, None, None] * m[None, :, None]) B = B * m[None, :, None] n_eff = jnp.maximum(jnp.sum(m), 1.0).astype(A.dtype) else: n_eff = jnp.asarray(float(n), dtype=A.dtype) P = jnp.einsum("ikc,kjc->ijc", A, B) P = P / jnp.sqrt(n_eff) p_ij = _inline_norm_forward( self.ln_path, P, pathway="even", kfac_structural_mask=full_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) cat = jnp.concatenate([pair_ij, p_ij], axis=-1) hidden = fused_silu( _tagged_dense( self.ffn_in.weight, self.ffn_in.bias, cat, tag_id=self.ffn_in._use_id, pathway="even", kfac_structural_mask=full_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) ) delta = _tagged_dense( self.ffn_out.weight, self.ffn_out.bias, hidden, tag_id=self.ffn_out._use_id, pathway="even", kfac_structural_mask=full_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) return delta def apply_residual( self, edge: Float[Array, "n n d_edge"], c_level: Float[Array, "n d_c"], mask: Float[Array, "n"] | None = None, *, kfac_scan_shared: bool = False, ) -> Float[Array, "n n d_edge"]: delta = self(edge, c_level, mask, kfac_scan_shared=kfac_scan_shared) update_mask = None if mask is not None: m = mask.astype(bool) update_mask = m[:, None] & m[None, :] return _tree_ngpt_residual( edge, delta, self.alpha, max_gain=self.ngpt_alpha_max, tag_id="", update_mask=update_mask, kfac_structural_mask=update_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) class CausalRouterEdgeFWLUpdate(EdgeFWLUpdate): two_hop_channels: int = eqx.field(static=True) def __init__( self, d_c: int, d_edge: int, *, key: PRNGKeyArray, alpha_init: float, alpha_max: float, channels: int = 64, edge_node_ctx_dim: int | None = None, ): super().__init__( d_c, d_edge, key=key, alpha_init=alpha_init, alpha_max=alpha_max, channels=channels, edge_node_ctx_dim=edge_node_ctx_dim, ) self.two_hop_channels = int(channels) def __call__( self, edge: Float[Array, "n n d_edge"], c_level: Float[Array, "n d_c"], mask: Float[Array, "n"] | None = None, *, kfac_scan_shared: bool = False, ) -> Float[Array, "n n d_edge"]: n = c_level.shape[0] node_structural_mask = mask full_pair_structural_mask = ( None if mask is None else mask[:, None] * mask[None, :] ) idx = jnp.arange(n, dtype=jnp.int32) causal_pair = idx[None, :] <= idx[:, None] if full_pair_structural_mask is None: full_pair_structural_mask = jnp.ones((n, n), dtype=bool) causal_pair_structural_mask = ( full_pair_structural_mask.astype(bool) & causal_pair ) edge_ln = _inline_norm_forward( self.ln_edge, edge, pathway="even", kfac_structural_mask=full_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) c_ln = _inline_norm_forward( self.ln_c, c_level, pathway="even", kfac_structural_mask=node_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) c_ctx = ( c_ln if self.node_ctx_proj is None else _tagged_dense_no_bias( self.node_ctx_proj.weight, c_ln, tag_id=self.node_ctx_proj._use_id, pathway="even", kfac_structural_mask=node_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) ) d_c_dim = c_ctx.shape[-1] c_i_b = jnp.broadcast_to(c_ctx[:, None, :], (n, n, d_c_dim)) c_j_b = jnp.broadcast_to(c_ctx[None, :, :], (n, n, d_c_dim)) pair_ij = jnp.concatenate([edge_ln, c_i_b, c_j_b], axis=-1) A = self._psi_apply( pair_ij, "L", kfac_structural_mask=causal_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, ) B = self._psi_apply( pair_ij, "R", kfac_structural_mask=full_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, ) if mask is not None: m = mask.astype(A.dtype) A = A * (m[:, None, None] * m[None, :, None]) B = B * m[None, :, None] allowed = causal_pair.astype(A.dtype) * m[None, :] else: allowed = causal_pair.astype(A.dtype) A = A * causal_pair[..., None].astype(A.dtype) n_eff = jnp.maximum(jnp.sum(allowed, axis=1), 1.0).astype(A.dtype) P = jnp.einsum("ikc,kjc->ijc", A, B) P = P / jnp.sqrt(n_eff)[:, None, None] P = P * causal_pair[..., None].astype(P.dtype) p_ij = _inline_norm_forward( self.ln_path, P, pathway="even", kfac_structural_mask=causal_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) cat = jnp.concatenate([pair_ij, p_ij], axis=-1) hidden = fused_silu( _tagged_dense( self.ffn_in.weight, self.ffn_in.bias, cat, tag_id=self.ffn_in._use_id, pathway="even", kfac_structural_mask=causal_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) ) delta = _tagged_dense( self.ffn_out.weight, self.ffn_out.bias, hidden, tag_id=self.ffn_out._use_id, pathway="even", kfac_structural_mask=causal_pair_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) delta_mask = causal_pair.astype(delta.dtype) if mask is not None: delta_mask = delta_mask * (m[:, None] * m[None, :]) return delta * delta_mask[..., None] def append_causal_row( self, edge: Float[Array, "n n d_edge"], c_level: Float[Array, "n d_c"], mask: Float[Array, "n"], row: Int[Array, ""], b_cache: Float[Array, "n n channels"], *, edge_row: Float[Array, "n d_edge"] | None = None, edge_col: Float[Array, "n d_edge"] | None = None, sequence_axis_name=None, sequence_mesh=None, ): n = c_level.shape[0] idx = jnp.arange(n, dtype=jnp.int32) row = jnp.asarray(row, dtype=jnp.int32) active = mask.astype(bool) allowed = active & (idx <= row) if edge_row is None: edge_row = edge[row] if edge_col is None: edge_col = edge[:, row] edge_row_ln = _inline_norm_forward( self.ln_edge, edge_row, pathway="even", kfac_structural_mask=allowed, kfac_repeat_ndim=1, ) edge_col_ln = _inline_norm_forward( self.ln_edge, edge_col, pathway="even", kfac_structural_mask=active, kfac_repeat_ndim=1, ) c_ln = _inline_norm_forward( self.ln_c, c_level, pathway="even", kfac_structural_mask=active, kfac_repeat_ndim=1, ) c_ctx = ( c_ln if self.node_ctx_proj is None else _tagged_dense_no_bias( self.node_ctx_proj.weight, c_ln, tag_id=self.node_ctx_proj._use_id, pathway="even", kfac_structural_mask=active, kfac_repeat_ndim=1, ) ) c_row = c_ctx[row] c_row_b = jnp.broadcast_to(c_row, c_ctx.shape) pair_row = jnp.concatenate([edge_row_ln, c_row_b, c_ctx], axis=-1) pair_col = jnp.concatenate([edge_col_ln, c_ctx, c_row_b], axis=-1) a_row = self._psi_apply( pair_row, "L", kfac_structural_mask=allowed, kfac_repeat_ndim=1 ) b_row = self._psi_apply( pair_row, "R", kfac_structural_mask=active, kfac_repeat_ndim=1 ) b_col = self._psi_apply( pair_col, "R", kfac_structural_mask=active, kfac_repeat_ndim=1 ) if sequence_axis_name is not None: from jax.sharding import NamedSharding, PartitionSpec as P lanes = ( int(sequence_mesh.shape[sequence_axis_name]) if sequence_mesh is not None else 1 ) row_spec = P(None, None) col_spec = P( sequence_axis_name if n >= lanes and n % lanes == 0 else None, None ) if sequence_mesh is not None: row_spec = NamedSharding(sequence_mesh, row_spec) col_spec = NamedSharding(sequence_mesh, col_spec) b_row = jax.lax.with_sharding_constraint(b_row, row_spec) b_col = jax.lax.with_sharding_constraint(b_col, col_spec) b_cache = _replace_square_row_column(b_cache, row, b_row, b_col) a_row = a_row * allowed[:, None].astype(a_row.dtype) path = jnp.einsum("kc,kjc->jc", a_row, b_cache) n_eff = jnp.maximum( jnp.sum(allowed.astype(path.dtype)), jnp.asarray(1.0, dtype=path.dtype) ) path = path / jnp.sqrt(n_eff) path = path * allowed[:, None].astype(path.dtype) path_ln = _inline_norm_forward( self.ln_path, path, pathway="even", kfac_structural_mask=allowed, kfac_repeat_ndim=1, ) cat = jnp.concatenate([pair_row, path_ln], axis=-1) hidden = fused_silu( _tagged_dense( self.ffn_in.weight, self.ffn_in.bias, cat, tag_id=self.ffn_in._use_id, pathway="even", kfac_structural_mask=allowed, kfac_repeat_ndim=1, ) ) delta = _tagged_dense( self.ffn_out.weight, self.ffn_out.bias, hidden, tag_id=self.ffn_out._use_id, pathway="even", kfac_structural_mask=allowed, kfac_repeat_ndim=1, ) delta = delta * allowed[:, None].astype(delta.dtype) updated = _tree_ngpt_residual( edge_row, delta, self.alpha, max_gain=self.ngpt_alpha_max, tag_id="", update_mask=allowed, kfac_structural_mask=allowed, kfac_repeat_ndim=1, ) return updated, b_cache def apply_residual( self, edge: Float[Array, "n n d_edge"], c_level: Float[Array, "n d_c"], mask: Float[Array, "n"] | None = None, *, kfac_scan_shared: bool = False, ) -> Float[Array, "n n d_edge"]: delta = self(edge, c_level, mask, kfac_scan_shared=kfac_scan_shared) idx = jnp.arange(edge.shape[0], dtype=jnp.int32) update_mask = idx[None, :] <= idx[:, None] if mask is not None: m = mask.astype(bool) update_mask = update_mask & m[:, None] & m[None, :] return _tree_ngpt_residual( edge, delta, self.alpha, max_gain=self.ngpt_alpha_max, tag_id="", update_mask=update_mask, kfac_structural_mask=update_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) class LevelEdgeAttn(eqx.Module): ln_scale: Float[Array, "d_c"] w_qkv: Float[Array, "d_c d_qkv"] w_o: Float[Array, "d_o_in d_c"] bias_mlp: MLP ffn_ln_scale: Float[Array, "d_c"] ffn_w1: Float[Array, "d_c d_ffn_hidden"] ffn_b1: Float[Array, "d_ffn_hidden"] ffn_w2: Float[Array, "d_ffn_hidden d_c"] ffn_b2: Float[Array, "d_c"] alpha_attn: Float[Array, "d_c"] alpha_ffn: Float[Array, "d_c"] n_heads: int = eqx.field(static=True) n_heads_kernel: int = eqx.field(static=True) d_attn: int = eqx.field(static=True) d_head: int = eqx.field(static=True) d_ffn_hidden: int = eqx.field(static=True) attn_impl: str = eqx.field(static=True) ln_eps: float = eqx.field(static=True) max_n: int = eqx.field(static=True) rope_base: float = eqx.field(static=True) rope_scaling: float = eqx.field(static=True) ngpt_alpha_max: float = eqx.field(static=True, default=0.5) _use_id_ln: str = eqx.field(static=True, default="") _use_id_qkv: str = eqx.field(static=True, default="") _use_id_o: str = eqx.field(static=True, default="") _use_id_ffn_ln: str = eqx.field(static=True, default="") _use_id_ffn1: str = eqx.field(static=True, default="") _use_id_ffn2: str = eqx.field(static=True, default="") def __init__( self, d_c: int, d_edge: int, *, key: PRNGKeyArray, alpha_init: float, alpha_max: float, n_heads: int = 4, attn_dim: int | None = None, attn_impl: str = "mhsea_tuned", bias_mlp_hidden: int | None = None, bias_mlp_n_blocks: int = 1, ffn_d_hidden: int | None = None, ln_eps: float = 1e-05, max_n: int = 128, rope_base: float = 10000.0, rope_scaling: float = 1.0, ): d_attn = int(d_c) if attn_dim is None else int(attn_dim) if d_attn < 1: raise ValueError( f"LevelEdgeAttn attn_dim must be positive or None, got {attn_dim}" ) assert d_attn % n_heads == 0, ( f"attn_dim ({d_attn}) must be divisible by n_heads ({n_heads})" ) if attn_impl not in ("einsum", "mhsea_tuned"): raise ValueError("attn_impl must be 'einsum' or 'mhsea_tuned'") n_heads_kernel = 2 * n_heads d_head = d_attn // n_heads if d_head % 2 != 0: raise ValueError( f"LevelEdgeAttn requires even d_head for RoPE, got d_head={d_head} (= attn_dim={d_attn} / n_heads={n_heads})" ) if max_n < 2: raise ValueError(f"LevelEdgeAttn max_n must be >= 2, got {max_n}") if rope_base <= 0.0: raise ValueError("LevelEdgeAttn rope_base must be positive") if rope_scaling <= 0.0: raise ValueError("LevelEdgeAttn rope_scaling must be positive") d_qkv_out = n_heads_kernel * d_head d_o_in = n_heads * d_head k_qkv, k_b, k_f1, k_o, k_f2 = jax.random.split(key, 5) self.w_qkv = jax.random.normal(k_qkv, (d_c, 3 * d_qkv_out)) * d_c ** (-0.5) self.w_o = jax.random.normal(k_o, (d_o_in, d_c)) * d_o_in ** (-0.5) if bias_mlp_hidden is None: bias_mlp_hidden = max(32, n_heads_kernel * 2) self.bias_mlp = MLP( d_edge, bias_mlp_hidden, n_heads_kernel, key=k_b, n_blocks=bias_mlp_n_blocks ) self.ln_scale = jnp.ones((d_c,)) d_ffn_eff = ffn_d_hidden if ffn_d_hidden is not None else 4 * d_c self.ffn_ln_scale = jnp.ones((d_c,)) self.ffn_w1 = jax.random.normal(k_f1, (d_c, d_ffn_eff)) * d_c ** (-0.5) self.ffn_b1 = jnp.zeros((d_ffn_eff,)) self.ffn_w2 = jax.random.normal(k_f2, (d_ffn_eff, d_c)) * d_ffn_eff ** (-0.5) self.ffn_b2 = jnp.zeros((d_c,)) self.n_heads = n_heads self.n_heads_kernel = n_heads_kernel self.d_attn = d_attn self.d_head = d_head self.d_ffn_hidden = d_ffn_eff self.attn_impl = attn_impl self.ln_eps = ln_eps self.max_n = int(max_n) self.rope_base = float(rope_base) self.rope_scaling = float(rope_scaling) self.ngpt_alpha_max = float(alpha_max) self.alpha_attn = float(alpha_init) * jnp.ones((int(d_c),)) self.alpha_ffn = float(alpha_init) * jnp.ones((int(d_c),)) def __call__( self, c_level: Float[Array, "n d_c"], edge: Float[Array, "n n d_edge"], mask: Float[Array, "n"], level_idx=None, *, kfac_scan_shared: bool = False, ) -> Float[Array, "n d_c"]: n = c_level.shape[0] H_k = self.n_heads_kernel d_h = self.d_head pair_mask = mask[:, None] * mask[None, :] x = _tagged_rms_eqx_style( self.ln_scale, c_level, eps=self.ln_eps, tag_id=self._use_id_ln, pathway="even", kfac_structural_mask=mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) qkv = _tagged_dense_no_bias( self.w_qkv, x, tag_id=self._use_id_qkv, pathway="even", kfac_structural_mask=mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) qkv = qkv.reshape(n, 3, H_k, d_h) Q = qkv[:, 0] K = qkv[:, 1] V = qkv[:, 2] bmlp = self.bias_mlp b = _tagged_dense( bmlp.in_proj.weight, bmlp.in_proj.bias, edge, tag_id=bmlp.in_proj._use_id, pathway="even", kfac_structural_mask=pair_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) for nrm, l1, l2 in zip(bmlp.block_norms, bmlp.block_l1s, bmlp.block_l2s): normed = _inline_norm_forward( nrm, b, pathway="even", kfac_structural_mask=pair_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) inner = _tagged_dense( l1.weight, l1.bias, normed, tag_id=l1._use_id, pathway="even", kfac_structural_mask=pair_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) inner_act = bmlp._act(inner) inner_out = _tagged_dense( l2.weight, l2.bias, inner_act, tag_id=l2._use_id, pathway="even", kfac_structural_mask=pair_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) b = b + bmlp.inner_gain * inner_out b_normed = _inline_norm_forward( bmlp.out_norm, b, pathway="even", kfac_structural_mask=pair_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) coup_bias = _tagged_dense( bmlp.out_proj.weight, bmlp.out_proj.bias, b_normed, tag_id=bmlp.out_proj._use_id, pathway="even", kfac_structural_mask=pair_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=2, ) coup_bias = coup_bias / jnp.sqrt(d_h) if level_idx is not None: _idx = jnp.arange(n, dtype=jnp.int32) _lca = lca_alibi_bias(_idx, _idx, lca_fixed_slopes(H_k, dtype=x.dtype)) coup_bias = coup_bias + jnp.transpose(_lca, (1, 2, 0)) impl = self.attn_impl from .pallas_attention import ( mhsea_tuned_edge_attention, reference_edge_attention, ) if impl == "einsum": out = reference_edge_attention(Q, K, V, coup_bias, mask) elif impl == "mhsea_tuned": d_head_padded = max(16, d_h) pad_amount = d_head_padded - d_h scale = jnp.sqrt(jnp.float32(d_head_padded / d_h)) pad_shape = (n, H_k, pad_amount) Q_p = jnp.concatenate([Q * scale, jnp.zeros(pad_shape, Q.dtype)], axis=-1) K_p = jnp.concatenate([K, jnp.zeros(pad_shape, K.dtype)], axis=-1) V_p = jnp.concatenate([V, jnp.zeros(pad_shape, V.dtype)], axis=-1) out = mhsea_tuned_edge_attention(Q_p, K_p, V_p, coup_bias, mask)[..., :d_h] else: raise ValueError("attn_impl must be 'einsum' or 'mhsea_tuned'") gate_heads = out[:, : self.n_heads, :] value_heads = out[:, self.n_heads :, :] out = jax.nn.sigmoid(gate_heads) * value_heads out_flat = out.reshape(n, self.n_heads * d_h) delta = _tagged_dense_no_bias( self.w_o, out_flat, tag_id=self._use_id_o, pathway="even", kfac_structural_mask=mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) mask_q = mask.reshape(-1, 1).astype(delta.dtype) proposal_attn = mask_q * delta c_attn = _tree_ngpt_residual( c_level, proposal_attn, self.alpha_attn, max_gain=self.ngpt_alpha_max, tag_id="", update_mask=mask, kfac_structural_mask=mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) x_ffn = _tagged_rms_eqx_style( self.ffn_ln_scale, c_attn, eps=self.ln_eps, tag_id=self._use_id_ffn_ln, pathway="even", kfac_structural_mask=mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) h = _tagged_dense( self.ffn_w1, self.ffn_b1, x_ffn, tag_id=self._use_id_ffn1, pathway="even", kfac_structural_mask=mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) h = fused_silu(h) delta_ffn = _tagged_dense( self.ffn_w2, self.ffn_b2, h, tag_id=self._use_id_ffn2, pathway="even", kfac_structural_mask=mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) proposal_ffn = mask_q * delta_ffn return _tree_ngpt_residual( c_attn, proposal_ffn, self.alpha_ffn, max_gain=self.ngpt_alpha_max, tag_id="", update_mask=mask, kfac_structural_mask=mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) class MergeOp(eqx.Module): T: Float[Array, "G d_r d_r d_r"] _use_id_T: str = eqx.field(static=True, default="") mlp_c: MLP alpha_c: Float[Array, "d_c"] _use_id_alpha_c: str = eqx.field(static=True, default="") ngpt_alpha_max: float = eqx.field(static=True, default=0.5) eps: float = eqx.field(static=True) output_hypernet: HypernetMatrix edge_merge: EdgeMergeOp level_edge_attn: LevelEdgeAttn tree_edge_fwl: EdgeFWLUpdate def __init__( self, d_r: int, d_c: int, *, key: PRNGKeyArray, d_g: int, alpha_init: float, alpha_max: float, eps: float = 1e-06, merge_output_hypernet_rank: int = 128, d_m_merge: int | None = None, level_edge_attn_d_edge: int = 64, level_edge_attn_n_heads: int = 4, level_edge_attn_attn_dim: int | None = None, tree_edge_node_ctx_dim: int | None = None, level_edge_attn_attn_impl: str = "mhsea_tuned", level_edge_attn_edge_mlp_hidden: int | None = None, level_edge_attn_edge_mlp_n_blocks: int = 2, level_edge_attn_ffn_d_hidden: int | None = None, level_edge_attn_max_n: int = 128, level_edge_attn_rope_base: float = 10000.0, level_edge_attn_rope_scaling: float = 1.0, tree_edge_fwl_channels: int = 64, level_edge_attn_bias_mlp_hidden: int | None = None, level_edge_attn_bias_mlp_n_blocks: int = 1, merge_c_mlp_hidden: int | None = None, ): keys = jax.random.split(key, 11) k_t, k_c, k_h = keys[:3] k_em, k_lea, k_fwl = keys[4], keys[5], keys[10] d_m_eff = d_m_merge if d_m_merge is not None else d_r if d_m_eff % d_r != 0: raise ValueError( f"MergeOp: d_m_merge={d_m_merge} must be divisible by d_r={d_r} (HT carrier T is reshape-indexed as [G, d_r, d_r, d_r] with G = d_m_eff // d_r)." ) G_merge = d_m_eff // d_r merge_edge_dim = 2 * int(level_edge_attn_d_edge) merge_clock_dim = int(d_c) merge_ngpt_dim = _TREE_NGPT_DEPTH_FEAT_DIM merge_extra_dim = int(d_g) + merge_edge_dim + merge_clock_dim + merge_ngpt_dim std = 1.0 / d_r self.T = jax.random.normal(k_t, (G_merge, d_r, d_r, d_r)) * std _c_hidden = ( int(merge_c_mlp_hidden) if merge_c_mlp_hidden is not None else max(d_c, 16) ) self.mlp_c = MLP(2 * d_c + merge_extra_dim, _c_hidden, d_c, key=k_c) self.ngpt_alpha_max = float(alpha_max) self.alpha_c = float(alpha_init) * jnp.ones((int(d_c),)) self.output_hypernet = HypernetMatrix( d_in=d_m_eff, d_out=d_m_eff, d_e=d_c + merge_ngpt_dim, rank=merge_output_hypernet_rank, key=k_h, ) self.eps = eps self.edge_merge = EdgeMergeOp( d_edge=level_edge_attn_d_edge, d_c=d_c, key=k_em, alpha_init=alpha_init, alpha_max=alpha_max, d_hidden=level_edge_attn_edge_mlp_hidden, n_blocks=level_edge_attn_edge_mlp_n_blocks, edge_node_ctx_dim=tree_edge_node_ctx_dim, ) self.level_edge_attn = LevelEdgeAttn( d_c=d_c, d_edge=level_edge_attn_d_edge, key=k_lea, alpha_init=alpha_init, alpha_max=alpha_max, n_heads=level_edge_attn_n_heads, attn_dim=level_edge_attn_attn_dim, attn_impl=level_edge_attn_attn_impl, ffn_d_hidden=level_edge_attn_ffn_d_hidden, max_n=level_edge_attn_max_n, rope_base=level_edge_attn_rope_base, rope_scaling=level_edge_attn_rope_scaling, bias_mlp_hidden=level_edge_attn_bias_mlp_hidden, bias_mlp_n_blocks=level_edge_attn_bias_mlp_n_blocks, ) self.tree_edge_fwl = EdgeFWLUpdate( d_c=d_c, d_edge=level_edge_attn_d_edge, key=k_fwl, alpha_init=alpha_init, alpha_max=alpha_max, channels=tree_edge_fwl_channels, edge_node_ctx_dim=tree_edge_node_ctx_dim, ) def _merge_extra_inputs( self, c_a, c_b, g_emb, sibling_edge_lr, sibling_edge_rl, level_idx, pair_idx, pair_base, clock_depth, depth_feats, ): parts = [g_emb] parts.append( jnp.concatenate( [sibling_edge_lr.astype(c_a.dtype), sibling_edge_rl.astype(c_a.dtype)] ) ) parts.append( _tree_merge_clock( level_idx, pair_idx, pair_base, c_a.shape[-1], clock_depth, c_a.dtype ) ) parts.append(depth_feats.astype(c_a.dtype)) return parts def _apply_c_skip( self, c_a, c_b, c_delta, *, kfac_structural_mask=None, kfac_scan_shared: bool = False, ): return _tree_ngpt_residual( 0.5 * (c_a + c_b), c_delta, self.alpha_c, max_gain=self.ngpt_alpha_max, tag_id=self._use_id_alpha_c, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) def context_candidate( self, c_a: Float[Array, "d_c"], c_b: Float[Array, "d_c"], g_emb: Float[Array, "d_g"], *, sibling_edge_lr: Float[Array, "d_edge"], sibling_edge_rl: Float[Array, "d_edge"], level_idx: Array, pair_idx: Array, pair_base: Array, clock_depth: Array, depth_feats: Array, kfac_structural_mask=None, kfac_g_structural_mask=None, kfac_scan_shared: bool = False, ): ffn_in = jnp.concatenate( [ c_a, c_b, *self._merge_extra_inputs( c_a, c_b, g_emb, sibling_edge_lr, sibling_edge_rl, level_idx, pair_idx, pair_base, clock_depth, depth_feats, ), ] ) mlp = self.mlp_c _we = _rownorm_cols x = _tagged_dense( mlp.in_proj.weight, mlp.in_proj.bias, ffn_in, tag_id=mlp.in_proj._use_id, pathway="even", weight_eff=_we(mlp.in_proj.weight), kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) for nrm, l1, l2 in zip(mlp.block_norms, mlp.block_l1s, mlp.block_l2s): normed = _inline_norm_forward( nrm, x, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) inner = _tagged_dense( l1.weight, l1.bias, normed, tag_id=l1._use_id, pathway="even", weight_eff=_we(l1.weight), kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) inner_act = mlp._act(inner) inner_out = _tagged_dense( l2.weight, l2.bias, inner_act, tag_id=l2._use_id, pathway="even", weight_eff=_we(l2.weight), kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) x = x + mlp.inner_gain * inner_out out_normed = _inline_norm_forward( mlp.out_norm, x, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) c_delta = _tagged_dense( mlp.out_proj.weight, mlp.out_proj.bias, out_normed, tag_id=mlp.out_proj._use_id, pathway="even", weight_eff=_we(mlp.out_proj.weight), kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) return self._apply_c_skip( c_a, c_b, c_delta, kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=kfac_scan_shared, ) def __call__( self, c_a: Float[Array, "d_c"], u_a: Float[Array, "d_m_eff"], s_a: Float[Array, ""], c_b: Float[Array, "d_c"], u_b: Float[Array, "d_m_eff"], s_b: Float[Array, ""], g_emb: Float[Array, "d_g"], *, sibling_edge_lr: Float[Array, "d_edge"], sibling_edge_rl: Float[Array, "d_edge"], level_idx: Array, pair_idx: Array, pair_base: Array, clock_depth: Array, depth_feats: Array, kfac_context_mask=None, kfac_g_context_mask=None, kfac_odd_mask=None, kfac_scan_shared: bool = False, ): c_p = self.context_candidate( c_a, c_b, g_emb, sibling_edge_lr=sibling_edge_lr, sibling_edge_rl=sibling_edge_rl, level_idx=level_idx, pair_idx=pair_idx, pair_base=pair_base, clock_depth=clock_depth, depth_feats=depth_feats, kfac_structural_mask=kfac_context_mask, kfac_g_structural_mask=kfac_g_context_mask, kfac_scan_shared=kfac_scan_shared, ) raw = _quadrilinear_merge( self.T, u_a, u_b, tag_id=self._use_id_T, pathway="odd", kfac_structural_mask=kfac_odd_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) H = self.output_hypernet h_ctx = jnp.concatenate([c_p, depth_feats.astype(c_p.dtype)]) h_p = _tagged_dense_no_bias( H.W_h, h_ctx, tag_id=H._use_id_W_h, pathway="hypernet_eside", kfac_structural_mask=kfac_odd_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, kfac_context_primal_reused_over_walkers=True, ) V_x = _tagged_dense_no_bias( H.V, raw, tag_id=H._use_id_V, pathway="odd", kfac_structural_mask=kfac_odd_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) m_p = h_p * V_x H_out = _tagged_dense_no_bias( H.U, m_p, tag_id=H._use_id_U, pathway="odd", kfac_structural_mask=kfac_odd_mask, kfac_scan_shared=kfac_scan_shared, kfac_repeat_ndim=1, ) out = raw + H_out scale_sq = jnp.mean(out * out) scale = jnp.sqrt(scale_sq + self.eps) u_p = out / scale s_p = s_a + s_b + jnp.log(scale) return (c_p, u_p, s_p) def project_tree_ngpt_rownorm(model): import equinox as _eqx def _collect(m): mlp = m.merge.mlp_c leaves = [mlp.in_proj.weight] for l1, l2 in zip(mlp.block_l1s, mlp.block_l2s): leaves.append(l1.weight) leaves.append(l2.weight) leaves.append(mlp.out_proj.weight) leaves.append(m.route_decoder.tree_merge.w1) leaves.append(m.route_decoder.tree_merge.w2) return tuple(leaves) targets = _collect(model) return _eqx.tree_at(_collect, model, _rownorm_project_jit(targets)) @partial(jax.jit, donate_argnums=0) def _rownorm_project_jit(ws): return tuple((_rownorm_cols(w) for w in ws)) def merge_masked( c_a, u_a, s_a, m_a, c_b, u_b, s_b, m_b, merge, g_emb, *, k_a, k_b, sibling_edge_lr, sibling_edge_rl, level_idx, pair_idx, pair_base, clock_depth, depth_feats, kfac_g_context_mask=None, kfac_scan_shared=False, ): both = m_a * m_b only_a = m_a * (1.0 - m_b) only_b = (1.0 - m_a) * m_b both_k = k_a * k_b only_a_k = k_a * (1.0 - k_b) only_b_k = (1.0 - k_a) * k_b def gate(value, left, right): return jnp.where( both.astype(bool), value, jnp.where( only_a.astype(bool), left, jnp.where(only_b.astype(bool), right, jnp.zeros_like(value)), ), ) def structural_gate(value, left, right): return jnp.where( both_k.astype(bool), value, jnp.where( only_a_k.astype(bool), left, jnp.where(only_b_k.astype(bool), right, jnp.zeros_like(value)), ), ) c_new, u_new, s_new = merge( c_a, u_a, s_a, c_b, u_b, s_b, g_emb, sibling_edge_lr=sibling_edge_lr, sibling_edge_rl=sibling_edge_rl, level_idx=level_idx, pair_idx=pair_idx, pair_base=pair_base, clock_depth=clock_depth, depth_feats=depth_feats, kfac_context_mask=both_k, kfac_g_context_mask=kfac_g_context_mask, kfac_odd_mask=both, kfac_scan_shared=kfac_scan_shared, ) c_out = structural_gate(c_new, c_a, c_b) out = (c_out, gate(u_new, u_a, u_b), gate(s_new, s_a, s_b), m_a + m_b - m_a * m_b) return out def edge_merge_masked( e_2i_2j: Float[Array, "d_edge"], e_2i_2j1: Float[Array, "d_edge"], e_2i1_2j: Float[Array, "d_edge"], e_2i1_2j1: Float[Array, "d_edge"], m_2i: Float[Array, ""], m_2i1: Float[Array, ""], m_2j: Float[Array, ""], m_2j1: Float[Array, ""], c_2i: Float[Array, "d_c"], c_2i1: Float[Array, "d_c"], c_2j: Float[Array, "d_c"], c_2j1: Float[Array, "d_c"], edge_merge: EdgeMergeOp, *, k_2i: Float[Array, ""], k_2i1: Float[Array, ""], k_2j: Float[Array, ""], k_2j1: Float[Array, ""], kfac_scan_shared: bool = False, ) -> tuple[Float[Array, "d_edge"], Float[Array, ""]]: m_p = m_2i + m_2i1 - m_2i * m_2i1 m_q = m_2j + m_2j1 - m_2j * m_2j1 m_pq = m_p * m_q both_p = k_2i * k_2i1 both_q = k_2j * k_2j1 out_mask = both_p * both_q proposal = edge_merge( e_2i_2j, e_2i_2j1, e_2i1_2j, e_2i1_2j1, c_2i, c_2i1, c_2j, c_2j1, kfac_structural_mask=out_mask, kfac_scan_shared=kfac_scan_shared, ) cell_weights = jnp.stack( [k_2i * k_2j, k_2i * k_2j1, k_2i1 * k_2j, k_2i1 * k_2j1] ).astype(proposal.dtype) child_edges = jnp.stack([e_2i_2j, e_2i_2j1, e_2i1_2j, e_2i1_2j1], axis=0) mean_denom = jnp.maximum( jnp.sum(cell_weights), jnp.asarray(1.0, dtype=proposal.dtype) ) masked_mean = jnp.sum(child_edges * cell_weights[:, None], axis=0) / mean_denom e_pq = edge_merge.apply_skip( masked_mean, proposal, kfac_structural_mask=out_mask, kfac_scan_shared=kfac_scan_shared, ) return ( jnp.where(jnp.asarray(out_mask).astype(bool), e_pq, jnp.zeros_like(e_pq)), m_pq, ) def _next_pow2(n: int) -> int: return 1 if n <= 1 else 1 << (n - 1).bit_length() def _balanced_subtree_mask(m, n_pad: int): N = jnp.sum(m.astype(jnp.int32)) powers = 2 ** jnp.arange(max(1, n_pad.bit_length()), dtype=jnp.int32) big = jnp.asarray(1 << 30, dtype=jnp.int32) next_p = jnp.min(jnp.where(powers >= jnp.maximum(N, 1), powers, big)) return (jnp.arange(n_pad, dtype=jnp.int32) < next_p).astype(m.dtype) def _tree_active_clock_depth(mask): n_active = jnp.maximum(jnp.sum(mask.astype(jnp.int32)), 1) depth = jnp.ceil(jnp.log2(n_active.astype(jnp.float32))).astype(jnp.int32) return jnp.maximum(depth, 1) def balanced_tree_reduce_masked_scan( c: Float[Array, "n d_c"], u: Float[Array, "n d_m_eff"], s: Float[Array, "n"], m: Float[Array, "n"], merge: MergeOp, g_emb: Float[Array, "d_g"], *, edges_init: Float[Array, "n n d_edge"], gladder, g_stream0, ): from hamiltonzero.model.fp32 import compute_dtype as _compute_dtype _cdtype = _compute_dtype() if c.dtype != _cdtype: c = c.astype(_cdtype) if edges_init.dtype != _cdtype: edges_init = edges_init.astype(_cdtype) n = c.shape[0] if n == 1: return (c[0], u[0], s[0], m[0], edges_init[0, 0], g_stream0) n_pad = _next_pow2(n) n_levels = n_pad.bit_length() - 1 pad_amount = n_pad - n if pad_amount > 0: c = jnp.pad(c, ((0, pad_amount),) + ((0, 0),) * (c.ndim - 1)) u = jnp.pad(u, ((0, pad_amount),) + ((0, 0),) * (u.ndim - 1)) s = jnp.pad(s, (0, pad_amount)) m = jnp.pad(m, (0, pad_amount)) k = _balanced_subtree_mask(m, n_pad) clock_depth = _tree_active_clock_depth(k) d_edge = edges_init.shape[-1] edges_padded = _tree_sphere( jnp.pad(edges_init, ((0, pad_amount), (0, pad_amount), (0, 0))) ) initial_state = (c, u, s, m, k, edges_padded, g_stream0) n_pairs = n_pad // 2 pidx = jnp.arange(n_pairs, dtype=jnp.int32) def _merge_one_pair( c_a_, u_a_, s_a_, c_b_, u_b_, s_b_, m_a_, m_b_, k_a_, k_b_, sibling_edge_lr_, sibling_edge_rl_, depth_feats_, pair_idx_, pair_base_, clock_depth_, level_idx_, g_emb_, g_structural_mask_, ): return merge_masked( c_a_, u_a_, s_a_, m_a_, c_b_, u_b_, s_b_, m_b_, merge, g_emb_, k_a=k_a_, k_b=k_b_, sibling_edge_lr=sibling_edge_lr_, sibling_edge_rl=sibling_edge_rl_, level_idx=level_idx_, pair_idx=pair_idx_, pair_base=pair_base_, clock_depth=clock_depth_, depth_feats=depth_feats_, kfac_g_context_mask=g_structural_mask_, kfac_scan_shared=True, ) _vmap_in_axes = (0,) * 14 + (None,) * 5 _vmapped_merge = jax.vmap( _merge_one_pair, in_axes=_vmap_in_axes, axis_name="tree_pair" ) def _edge_one( e0, e1, e2, e3, m_p_a, m_p_b, m_q_a, m_q_b, k_p_a, k_p_b, k_q_a, k_q_b, c_p_a, c_p_b, c_q_a, c_q_b, ): return edge_merge_masked( e0, e1, e2, e3, m_p_a, m_p_b, m_q_a, m_q_b, c_p_a, c_p_b, c_q_a, c_q_b, merge.edge_merge, k_2i=k_p_a, k_2i1=k_p_b, k_2j=k_q_a, k_2j1=k_q_b, kfac_scan_shared=True, ) _edge_inner = jax.vmap( _edge_one, in_axes=(0, 0, 0, 0, None, None, 0, 0, None, None, 0, 0, None, None, 0, 0), axis_name="tree_edge_q", ) _edge_outer = jax.vmap( _edge_inner, in_axes=(0, 0, 0, 0, 0, 0, None, None, 0, 0, None, None, 0, 0, None, None), axis_name="tree_edge_p", ) def body(state, xs_lv): level_idx, depth_feats_lv = xs_lv c, u, s, m, k, E, g_carry = state def _split(x): xr = x.reshape((n_pairs, 2) + x.shape[1:]) return (xr[:, 0], xr[:, 1]) m_a, m_b = _split(m) k_a, k_b = _split(k) level_active = jnp.any((k_a * k_b).astype(bool)) c_a, c_b = _split(c) u_a, u_b = _split(u) s_a, s_b = _split(s) pair_args: list = [c_a, u_a, s_a, c_b, u_b, s_b] pair_args.extend([m_a, m_b]) pair_args.extend([k_a, k_b]) E_rs_for_merge = E.reshape(n_pairs, 2, n_pairs, 2, d_edge) pair_args.extend( [E_rs_for_merge[pidx, 0, pidx, 1, :], E_rs_for_merge[pidx, 1, pidx, 0, :]] ) g_emb_lvl = _tagged_dense( gladder[2], gladder[3], g_carry, tag_id="gladder.tree.proj", pathway="even", kfac_structural_mask=level_active, kfac_scan_shared=True, kfac_repeat_ndim=0, ) pair_args.append(depth_feats_lv) clock_pair_active = k_a + k_b - k_a * k_b pair_base = jnp.maximum( jnp.sum(clock_pair_active.astype(jnp.int32)), jnp.asarray(2, dtype=jnp.int32), ) pair_args.extend( [pidx, pair_base, clock_depth, level_idx, g_emb_lvl, level_active] ) merged = _vmapped_merge(*pair_args) c_p, u_p, s_p, m_p = merged k_p = k_a + k_b - k_a * k_b both_struct = k_a * k_b attn_mask = both_struct E_rs = E.reshape(n_pairs, 2, n_pairs, 2, d_edge) E_00 = E_rs[:, 0, :, 0, :] E_01 = E_rs[:, 0, :, 1, :] E_10 = E_rs[:, 1, :, 0, :] E_11 = E_rs[:, 1, :, 1, :] E_new, _m_edge_new = _edge_outer( E_00, E_01, E_10, E_11, m_a, m_b, m_a, m_b, k_a, k_b, k_a, k_b, c_a, c_b, c_a, c_b, ) E_new = merge.tree_edge_fwl.apply_residual( E_new, c_p, attn_mask, kfac_scan_shared=True ) edge_keep = (both_struct[:, None] * both_struct[None, :]).astype(bool) E_new = jnp.where(edge_keep[..., None], E_new, E_00) E_new = jnp.where(edge_keep[..., None], _tree_sphere(E_new), E_00) c_skip = c_p c_p = merge.level_edge_attn( c_p, E_new, attn_mask, level_idx=level_idx, kfac_scan_shared=True ) c_p = jnp.where(attn_mask.astype(bool)[:, None], _tree_sphere(c_p), c_skip) _lvl_mask = k_p.astype(c_p.dtype) update_active = jnp.any(attn_mask.astype(bool)) pool_structural_mask = _lvl_mask * update_active.astype(_lvl_mask.dtype) pooled = gladder[0]( g_carry, c_p, _lvl_mask, kfac_structural_mask=pool_structural_mask, kfac_update_mask=update_active, kfac_scan_shared=True, kfac_repeat_ndim=1, ) g_carry = gladder[1]( g_carry, pooled, update_mask=update_active, kfac_structural_mask=update_active, kfac_scan_shared=True, ) def _zpad(x_half): return jnp.concatenate([x_half, jnp.zeros_like(x_half)], axis=0) pad_amt = n_pad - n_pairs E_padded = jnp.pad(E_new, ((0, pad_amt), (0, pad_amt), (0, 0))) return ( ( _zpad(c_p), _zpad(u_p), _zpad(s_p), _zpad(m_p), _zpad(k_p), E_padded, g_carry, ), None, ) depth_feat_levels = _tree_ngpt_level_counts(m, n_pairs, n_levels, c.dtype) final_state, _ = jax.lax.scan( body, initial_state, (jnp.arange(n_levels), depth_feat_levels) ) c_f, u_f, s_f, m_f, _k_f, E_f, g_final = final_state return (c_f[0], u_f[0], s_f[0], m_f[0], E_f[0, 0], g_final) class RootReadout(eqx.Module): output_hypernet: HypernetMatrix ln_e: _RMS def __init__( self, d_r: int, *, key: PRNGKeyArray, d_m_merge: int | None = None, d_edge: int, edge_rank: int = 64, d_g: int = 0, d_c: int, ): d_m_eff = d_m_merge if d_m_merge is not None else d_r if d_m_eff % d_r != 0: raise ValueError( f"RootReadout: d_m_merge={d_m_merge} must be divisible by d_r={d_r} (HT carrier dim must match MergeOp's d_m_eff)." ) keys = jax.random.split(key, 5) d_e_ctx = int(d_edge) + int(d_c) + int(d_g) self.output_hypernet = HypernetMatrix( d_in=d_m_eff, d_out=2, d_e=d_e_ctx, rank=edge_rank, key=keys[0] ) self.ln_e = _RMS(d_edge) def __call__(self, u_r, s_r, *, e_root, g_emb, c_root, kfac_structural_mask=None): e_norm = _inline_norm_forward( self.ln_e, e_root, pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=False, kfac_repeat_ndim=0, ) h_ctx = jnp.concatenate( [e_norm, c_root.astype(e_norm.dtype), g_emb.astype(e_norm.dtype)] ) psi = self.output_hypernet.apply( h_ctx, u_r, e_pathway="even", kfac_structural_mask=kfac_structural_mask, kfac_scan_shared=False, kfac_repeat_ndim=0, kfac_context_primal_reused_over_walkers=True, ) re = 0.5 * jnp.log(psi[0] * psi[0] + psi[1] * psi[1]) + s_r return (re, jnp.arctan2(psi[1], psi[0]))