File size: 899 Bytes
5ccb4fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import equinox as eqx
from jaxtyping import Array
class RouterKernel(eqx.Module):
contextualizer: object
global_fork: object
decoder: object
class RouterStatic(eqx.Module):
node_input: Array
node_projected: Array
global_input: Array
global_projected: Array
raw_edge: Array
initial_suffix: Array
prefix_edge_messages: Array
suffix_edge_messages: Array
order_decay: Array
virtual_decay: Array
tree_pair_messages: Array
static_bias_tables: tuple[Array, ...]
quadratic_base_static: tuple[Array, ...]
forked_global_static: tuple[Array, ...]
quotient_node_key: Array
quotient_edge_key: Array
real_mask: Array
routable_mask: Array
needs_fwl2: Array
__all__ = ["RouterKernel", "RouterStatic"]
|