| |
| |
|
|
| from __future__ import annotations |
|
|
| import jax |
| import jax.numpy as jnp |
|
|
| from hamiltonzero.compiled.types import SharedTrunk |
| from hamiltonzero.model.route_pointer import TreePrefixPointerMHSEA, lca_gaussian_decay |
|
|
| from .types import RouterKernel, RouterStatic |
|
|
|
|
| def bind_router_kernel(model) -> RouterKernel: |
| decoder = model.route_decoder |
| if not isinstance(decoder, TreePrefixPointerMHSEA): |
| raise ValueError("learned-router train requires TreePrefixPointerMHSEA") |
| return RouterKernel( |
| contextualizer=model.route_contextualizer, |
| global_fork=model.gladder_fork_route, |
| decoder=decoder, |
| ) |
|
|
|
|
| def compile_router_static( |
| kernel: RouterKernel, |
| trunk: SharedTrunk, |
| quotient_node_key, |
| quotient_edge_key, |
| needs_fwl2, |
| ) -> RouterStatic: |
| decoder = kernel.decoder |
| node_input, raw_edge, global_input = kernel.contextualizer.with_edge( |
| trunk.node_raw, |
| trunk.edge_raw, |
| trunk.real_mask, |
| trunk.balanced_mask, |
| g=trunk.global_stream, |
| ) |
| global_input = kernel.global_fork(global_input, raw_edge, trunk.balanced_mask) |
| (node_raw, node_projected), _ = decoder._prepare_nodes( |
| node_input, trunk.balanced_mask |
| ) |
| global_raw, global_projected = decoder._project_global( |
| global_input, node_input.dtype |
| ) |
| ((_, _, initial_suffix, _, _), edge_messages) = ( |
| decoder._initial_summaries_and_edge_messages( |
| raw_edge, trunk.balanced_mask, node_input.dtype |
| ) |
| ) |
| prefix_edge_messages, suffix_edge_messages = edge_messages |
| n = node_input.shape[0] |
| idx = jnp.arange(n, dtype=jnp.int32) |
| return RouterStatic( |
| node_input=node_raw, |
| node_projected=node_projected, |
| global_input=global_raw, |
| global_projected=global_projected, |
| raw_edge=raw_edge, |
| initial_suffix=initial_suffix, |
| prefix_edge_messages=prefix_edge_messages, |
| suffix_edge_messages=suffix_edge_messages, |
| order_decay=lca_gaussian_decay( |
| idx, idx, decoder.order_decay_w[0], decoder.order_decay_b[0] |
| ), |
| virtual_decay=lca_gaussian_decay( |
| idx, idx, decoder.virt_decay_w[0], decoder.virt_decay_b[0] |
| ), |
| tree_pair_messages=decoder._tree_pair_messages(raw_edge, trunk.balanced_mask), |
| static_bias_tables=decoder._pack_heavy_static_bias_tables(raw_edge), |
| quadratic_base_static=(), |
| forked_global_static=(), |
| quotient_node_key=quotient_node_key, |
| quotient_edge_key=quotient_edge_key, |
| real_mask=trunk.real_mask, |
| routable_mask=trunk.balanced_mask, |
| needs_fwl2=jnp.asarray(needs_fwl2, dtype=jnp.bool_), |
| ) |
|
|
|
|
| __all__ = ["bind_router_kernel", "compile_router_static"] |
|
|
