File size: 3,559 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | # Copyright (c) 2026 Simulacra Research Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import functools
import jax
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
ROUTE_SAMPLES = 8
GLOBAL_BEAM_WIDTH = 16
def _decode_one(decoder, static, key, tau):
return decoder._decode(
static.node_input,
static.raw_edge,
static.routable_mask,
tau=tau,
key=key,
real_mask=static.real_mask,
first_orbit_ids=(
static.quotient_node_key,
static.quotient_edge_key,
static.needs_fwl2,
),
router_static=static,
)
def _beam_local(decoder, static, tau, *, lanes):
permutations, _log_probabilities = decoder.beam_search(
static.node_input,
static.raw_edge,
static.routable_mask,
global_feat=static.global_input,
tau=tau,
beam_width=GLOBAL_BEAM_WIDTH,
real_mask=static.real_mask,
first_orbit_ids=(
static.quotient_node_key,
static.quotient_edge_key,
static.needs_fwl2,
),
router_static=static,
distributed_axis_name="systems",
distributed_lanes=lanes,
)
return permutations[0]
def build_beam16(mesh: Mesh, decoder, static):
lanes = int(mesh.shape["systems"])
if tuple(mesh.axis_names) != ("systems",) or GLOBAL_BEAM_WIDTH % lanes:
raise ValueError("beam16 requires a one-dimensional divisible systems mesh")
mapped = jax.shard_map(
functools.partial(_beam_local, lanes=lanes),
mesh=mesh,
in_specs=(
jax.tree_util.tree_map(lambda _: P(), decoder),
jax.tree_util.tree_map(lambda _: P(), static),
P(),
),
out_specs=P(),
check_vma=False,
)
replicated = NamedSharding(mesh, P())
return jax.jit(
mapped,
in_shardings=(
jax.tree_util.tree_map(lambda _: replicated, decoder),
jax.tree_util.tree_map(lambda _: replicated, static),
replicated,
),
out_shardings=replicated,
)
def build_route_sampler(mesh: Mesh, decoder, static):
if tuple(mesh.axis_names) != ("systems",) or mesh.shape["systems"] != ROUTE_SAMPLES:
raise ValueError("learned-router train requires an eight-lane systems mesh")
replicated = NamedSharding(mesh, P())
route_vector = NamedSharding(mesh, P("systems", None))
local_specs = (
jax.tree_util.tree_map(lambda _: P(), decoder),
jax.tree_util.tree_map(lambda _: P(), static),
P(),
P(),
)
def local(decoder_value, static_value, key, tau):
lane_key = jax.random.fold_in(key, jax.lax.axis_index("systems"))
sample_key = jax.random.split(lane_key, 1)[0]
permutation = _decode_one(
decoder_value,
static_value,
sample_key,
tau,
)
return permutation[None]
mapped = jax.shard_map(
local,
mesh=mesh,
in_specs=local_specs,
out_specs=P("systems", None),
check_vma=False,
)
return jax.jit(
mapped,
in_shardings=(
jax.tree_util.tree_map(lambda _: replicated, decoder),
jax.tree_util.tree_map(lambda _: replicated, static),
replicated,
replicated,
),
out_shardings=route_vector,
)
__all__ = [
"GLOBAL_BEAM_WIDTH",
"ROUTE_SAMPLES",
"build_beam16",
"build_route_sampler",
]
|