"""Tests for the mixture-of-experts routing prototype (giant/model/network.py).""" import torch from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( ComposedRouter, DenoisingMLP, EnergyRouter, PdgRouter, ProcessRouter, ROUTER_REGISTRY, RoutedDenoisingMLP, RoutedSecondaryDecoder, SecondaryDecoder, build_composed_router, build_models, build_router, ) def _cond(B=8, pdg=3, mat=2): cond_cont = torch.randn(B, COND_DIM) cond_cat = torch.stack( [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1 ) return cond_cont, cond_cat def _routed_stage1(n_experts=4, pdg=3, mat=2, **router_kwargs): router = build_router("energy", n_experts, **router_kwargs) return RoutedDenoisingMLP( pdg_vocab=pdg, mat_vocab=mat, router=router, expert_hidden_dim=16, expert_n_blocks=2, ) def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs): router = build_router("energy", n_experts, **router_kwargs) return RoutedSecondaryDecoder( pdg_vocab=pdg, mat_vocab=mat, router=router, expert_hidden_dim=16, expert_n_blocks=2, ) # ── Router / EnergyRouter contract ────────────────────────────────────────── def test_energy_router_registered(): assert ROUTER_REGISTRY["energy"] is EnergyRouter def test_energy_router_gate_partition_of_unity(): router = EnergyRouter(n_experts=4) cond_cont, cond_cat = _cond(16) g = router.gate(cond_cont, cond_cat) assert g.shape == (16, 4) torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) def test_energy_router_top1_matches_gate_argmax(): router = EnergyRouter(n_experts=4) cond_cont, cond_cat = _cond(16) assert torch.equal( router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1) ) def test_energy_router_hardens_as_temperature_shrinks(): """As tau -> 0 the soft gate should converge to a one-hot at the argmax.""" router = EnergyRouter(n_experts=4, temperature=1e-4) cond_cont, cond_cat = _cond(16) g = router.gate(cond_cont, cond_cat) top1 = router.top1(cond_cont, cond_cat) onehot = torch.nn.functional.one_hot(top1, num_classes=4).float() torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0) def test_energy_router_balance_loss_is_nonnegative_scalar(): router = EnergyRouter(n_experts=4) cond_cont, cond_cat = _cond(16) loss = router.balance_loss(cond_cont, cond_cat) assert loss.shape == () assert loss.item() >= 0.0 def test_build_router_ignores_unrecognized_kwargs(): # lambda_balance is a model_config.router key but not an EnergyRouter kwarg router = build_router("energy", 4, temperature=0.3, lambda_balance=0.5) assert isinstance(router, EnergyRouter) assert router.temperature == 0.3 def test_energy_router_default_centers_are_linspace(): router = EnergyRouter(n_experts=4) torch.testing.assert_close(router.centers, torch.linspace(-2.0, 2.0, 4)) def test_energy_router_centers_init_overrides_default(): centers_init = [-1.0, 0.0, 0.5, 3.0] router = EnergyRouter(n_experts=4, centers_init=centers_init) torch.testing.assert_close(router.centers, torch.tensor(centers_init)) def test_energy_router_centers_init_wrong_length_raises(): try: EnergyRouter(n_experts=4, centers_init=[0.0, 1.0]) except ValueError: return raise AssertionError("expected ValueError for centers_init length mismatch") def test_energy_router_centers_init_respects_learn_centers_flag(): learned = EnergyRouter( n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True ) fixed = EnergyRouter( n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False ) assert isinstance(learned.centers, torch.nn.Parameter) assert not isinstance(fixed.centers, torch.nn.Parameter) def test_build_router_threads_centers_init_through_energy_router(): centers_init = [-1.5, -0.5, 0.5, 1.5] router = build_router("energy", 4, centers_init=centers_init) torch.testing.assert_close(router.centers, torch.tensor(centers_init)) def test_build_router_unknown_type_raises(): try: build_router("nonexistent", 4) except ValueError: return raise AssertionError("expected ValueError for unknown router type") # ── PdgRouter ──────────────────────────────────────────────────────────────── def test_pdg_router_registered(): assert ROUTER_REGISTRY["pdg"] is PdgRouter def test_pdg_router_gate_partition_of_unity(): router = PdgRouter(n_experts=4, pdg_vocab=3) cond_cont, cond_cat = _cond(16) g = router.gate(cond_cont, cond_cat) assert g.shape == (16, 4) torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) def test_pdg_router_top1_matches_gate_argmax(): router = PdgRouter(n_experts=4, pdg_vocab=3) cond_cont, cond_cat = _cond(16) assert torch.equal( router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1) ) def test_pdg_router_hardens_as_temperature_shrinks(): """As tau -> 0 the soft gate should converge to a one-hot at the argmax.""" router = PdgRouter(n_experts=4, pdg_vocab=3, temperature=1e-4) cond_cont, cond_cat = _cond(16) g = router.gate(cond_cont, cond_cat) top1 = router.top1(cond_cont, cond_cat) onehot = torch.nn.functional.one_hot(top1, num_classes=4).float() torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0) def test_pdg_router_balance_loss_is_nonnegative_scalar(): router = PdgRouter(n_experts=4, pdg_vocab=3) cond_cont, cond_cat = _cond(16) loss = router.balance_loss(cond_cont, cond_cat) assert loss.shape == () assert loss.item() >= 0.0 def test_pdg_router_classify_loss_defaults_to_zero(): """PDG is already known at gate time (unlike ProcessRouter's process label), so no supervision is needed — falls back to Router's default.""" router = PdgRouter(n_experts=4, pdg_vocab=3) cond_cont, cond_cat = _cond(16) labels = torch.randint(0, 4, (16,)) loss = router.classify_loss(cond_cont, cond_cat, labels) assert loss.shape == () assert loss.item() == 0.0 def test_pdg_router_only_reads_pdg_column(): """Gate must depend on cond_cat[:, 0] (pdg) only, not cond_cont or material.""" router = PdgRouter(n_experts=4, pdg_vocab=3) cond_cont, cond_cat = _cond(16) g_before = router.gate(cond_cont, cond_cat) cond_cont_perturbed = torch.randn_like(cond_cont) cond_cat_diff_mat = cond_cat.clone() cond_cat_diff_mat[:, 1] = (cond_cat_diff_mat[:, 1] + 1) % 2 g_after = router.gate(cond_cont_perturbed, cond_cat_diff_mat) torch.testing.assert_close(g_before, g_after, atol=1e-6, rtol=0) def test_build_router_pdg_type_uses_pdg_vocab(): router = build_router("pdg", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8) assert isinstance(router, PdgRouter) assert router.pdg_emb.num_embeddings == 5 def test_build_models_routed_with_pdg_router(): model_config = dict( pdg_vocab=4, mat_vocab=2, emb_dim=16, dropout=0.1, k_max=K_MAX, expert_hidden_dim=16, expert_n_blocks=2, router={ "enabled": True, "type": "pdg", "n_experts": 3, }, ) stage1, sec_decoder = build_models(model_config) assert isinstance(stage1, RoutedDenoisingMLP) assert isinstance(stage1.router, PdgRouter) assert len(stage1.experts) == 3 assert stage1.router.pdg_emb.num_embeddings == 4 # ── ProcessRouter ──────────────────────────────────────────────────────────── def test_process_router_registered(): assert ROUTER_REGISTRY["process"] is ProcessRouter def test_process_router_gate_partition_of_unity(): router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) cond_cont, cond_cat = _cond(16) g = router.gate(cond_cont, cond_cat) assert g.shape == (16, 4) torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) def test_process_router_top1_matches_gate_argmax(): router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) cond_cont, cond_cat = _cond(16) assert torch.equal( router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1) ) def test_process_router_balance_loss_is_nonnegative_scalar(): router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) cond_cont, cond_cat = _cond(16) loss = router.balance_loss(cond_cont, cond_cat) assert loss.shape == () assert loss.item() >= 0.0 def test_process_router_classify_loss_decreases_with_training(): """The classifier should be able to fit an arbitrary label assignment — a sanity check that gradients actually flow to the process classifier.""" torch.manual_seed(0) router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2) cond_cont, cond_cat = _cond(32) labels = torch.randint(0, 4, (32,)) opt = torch.optim.Adam(router.parameters(), lr=0.05) first = router.classify_loss(cond_cont, cond_cat, labels).item() for _ in range(50): opt.zero_grad() loss = router.classify_loss(cond_cont, cond_cat, labels) loss.backward() opt.step() last = loss.item() assert last < first def test_energy_router_classify_loss_defaults_to_zero(): """Routers with no supervised signal (EnergyRouter) fall back to the Router base class's zero-loss default.""" router = EnergyRouter(n_experts=4) cond_cont, cond_cat = _cond(16) labels = torch.randint(0, 4, (16,)) loss = router.classify_loss(cond_cont, cond_cat, labels) assert loss.shape == () assert loss.item() == 0.0 def test_build_router_process_type_uses_pdg_mat_vocab(): router = build_router("process", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8) assert isinstance(router, ProcessRouter) assert router.pdg_emb.num_embeddings == 5 assert router.mat_emb.num_embeddings == 3 def test_build_models_routed_with_process_router(): model_config = dict( pdg_vocab=4, mat_vocab=2, emb_dim=16, dropout=0.1, k_max=K_MAX, expert_hidden_dim=16, expert_n_blocks=2, router={ "enabled": True, "type": "process", "n_experts": 3, "lambda_proc": 1.0, }, ) stage1, sec_decoder = build_models(model_config) assert isinstance(stage1, RoutedDenoisingMLP) assert isinstance(stage1.router, ProcessRouter) assert len(stage1.experts) == 3 assert stage1.router.pdg_emb.num_embeddings == 4 assert stage1.router.mat_emb.num_embeddings == 2 # ── ComposedRouter ─────────────────────────────────────────────────────────── def test_composed_router_n_experts_is_product(): router = ComposedRouter( [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)] ) assert router.n_experts == 12 def test_composed_router_gate_partition_of_unity(): router = ComposedRouter( [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)] ) cond_cont, cond_cat = _cond(16, pdg=5) g = router.gate(cond_cont, cond_cat) assert g.shape == (16, 12) torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) def test_composed_router_gate_is_outer_product_of_sub_gates(): energy_router = EnergyRouter(n_experts=4) pdg_router = PdgRouter(n_experts=3, pdg_vocab=5) router = ComposedRouter([energy_router, pdg_router]) cond_cont, cond_cat = _cond(16, pdg=5) g_energy = energy_router.gate(cond_cont, cond_cat) # (16, 4) g_pdg = pdg_router.gate(cond_cont, cond_cat) # (16, 3) expected = (g_energy.unsqueeze(-1) * g_pdg.unsqueeze(1)).flatten(1) # (16, 12) torch.testing.assert_close(router.gate(cond_cont, cond_cat), expected) def test_composed_router_top1_factors_into_per_axis_argmax(): """Joint argmax over the outer product must equal the pair of per-axis argmaxes, flattened with the same row-major index convention as gate().""" energy_router = EnergyRouter(n_experts=4) pdg_router = PdgRouter(n_experts=3, pdg_vocab=5) router = ComposedRouter([energy_router, pdg_router]) cond_cont, cond_cat = _cond(16, pdg=5) joint_idx = router.top1(cond_cont, cond_cat) energy_idx = energy_router.top1(cond_cont, cond_cat) pdg_idx = pdg_router.top1(cond_cont, cond_cat) expected = energy_idx * pdg_router.n_experts + pdg_idx assert torch.equal(joint_idx, expected) def test_composed_router_supports_different_expert_counts_per_axis(): router = ComposedRouter( [EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)] ) assert router.n_experts == 10 cond_cont, cond_cat = _cond(8, pdg=5) assert router.gate(cond_cont, cond_cat).shape == (8, 10) def test_composed_router_classify_loss_sums_sub_router_losses(): """energy/pdg both default to zero, so the composed loss should too.""" router = ComposedRouter( [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)] ) cond_cont, cond_cat = _cond(16, pdg=5) labels = torch.randint(0, 4, (16,)) loss = router.classify_loss(cond_cont, cond_cat, labels) assert loss.shape == () assert loss.item() == 0.0 def test_composed_router_rejects_empty_router_list(): try: ComposedRouter([]) except ValueError: return raise AssertionError("expected ValueError for empty router list") def test_composed_router_not_in_registry(): assert "composed" not in ROUTER_REGISTRY # ── _parse_composed_axes (axis{i}_{field} flat-key config convention) ─────── def test_parse_composed_axes_groups_indexed_keys(): from giant.model.network import _parse_composed_axes router_cfg = { "enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4, "axis0_temperature": 0.3, "axis1_type": "pdg", "axis1_n_experts": 3, "axis1_emb_dim": 6, } axes = _parse_composed_axes(router_cfg) assert axes == [ {"type": "energy", "n_experts": 4, "temperature": 0.3}, {"type": "pdg", "n_experts": 3, "emb_dim": 6}, ] def test_parse_composed_axes_ignores_unrelated_keys(): from giant.model.network import _parse_composed_axes router_cfg = { "enabled": True, "type": "composed", "lambda_balance": 0.0, "axis0_type": "energy", "axis0_n_experts": 4, } axes = _parse_composed_axes(router_cfg) assert axes == [{"type": "energy", "n_experts": 4}] def test_parse_composed_axes_raises_on_index_gap(): from giant.model.network import _parse_composed_axes router_cfg = { "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4, # axis1 missing entirely "axis2_type": "pdg", "axis2_n_experts": 3, } try: _parse_composed_axes(router_cfg) except ValueError: return raise AssertionError("expected ValueError for a gap in axis indices") def test_build_composed_router_resolves_per_axis_specs(): router = build_composed_router( [ {"type": "energy", "n_experts": 4, "temperature": 0.3}, {"type": "pdg", "n_experts": 3, "emb_dim": 6}, ], pdg_vocab=5, mat_vocab=2, ) assert isinstance(router, ComposedRouter) assert router.n_experts == 12 energy_router, pdg_router = router.routers assert isinstance(energy_router, EnergyRouter) assert energy_router.temperature == 0.3 assert isinstance(pdg_router, PdgRouter) assert pdg_router.pdg_emb.num_embeddings == 5 assert pdg_router.pdg_emb.embedding_dim == 6 def test_build_models_routed_with_composed_router(): model_config = dict( pdg_vocab=5, mat_vocab=2, emb_dim=16, dropout=0.1, k_max=K_MAX, expert_hidden_dim=16, expert_n_blocks=2, router={ "enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4, "axis1_type": "pdg", "axis1_n_experts": 3, }, ) stage1, sec_decoder = build_models(model_config) assert isinstance(stage1, RoutedDenoisingMLP) assert isinstance(stage1.router, ComposedRouter) assert len(stage1.experts) == 12 assert len(sec_decoder.experts) == 12 # stage1 and sec_decoder must not share router weights (same convention # as the single-axis routers built by build_models). assert stage1.router is not sec_decoder.router def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow(): from giant.sample import sample_flow, sample_secondaries model_config = dict( pdg_vocab=3, mat_vocab=2, emb_dim=16, dropout=0.1, k_max=K_MAX, expert_hidden_dim=8, expert_n_blocks=1, router={ "enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 2, "axis1_type": "pdg", "axis1_n_experts": 2, }, ) stage1, sec_decoder = build_models(model_config) B = 5 cond_cont, cond_cat = _cond(B, pdg=3, mat=2) stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2) assert stage1_norm.shape == (B, X_DIM) assert n_sec_pred.shape == (B,) sec_cont, sec_type_emb, sec_valid = sample_secondaries( sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2 ) assert sec_cont.shape == (B, K_MAX, 4) assert sec_valid.shape == (B, K_MAX) # ── RoutedDenoisingMLP ─────────────────────────────────────────────────────── def test_routed_denoising_mlp_output_shape_train_and_eval(): B = 8 model = _routed_stage1() x_t = torch.randn(B, X_DIM) t = torch.rand(B) cond_cont, cond_cat = _cond(B) model.train() out_train = model(x_t, t, cond_cont, cond_cat) assert out_train.shape == (B, X_DIM) model.eval() with torch.no_grad(): out_eval = model(x_t, t, cond_cont, cond_cat) assert out_eval.shape == (B, X_DIM) def test_routed_denoising_mlp_gradients_flow_in_train_mode(): """Soft mixture in train mode should touch every expert's parameters.""" B = 8 model = _routed_stage1(n_experts=3) x_t = torch.randn(B, X_DIM) t = torch.rand(B) cond_cont, cond_cat = _cond(B) model.train() flow_loss = model(x_t, t, cond_cont, cond_cat).sum() nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum() (flow_loss + nsec_loss).backward() for name, p in model.named_parameters(): assert p.grad is not None, f"no grad for {name}" def test_routed_denoising_mlp_eval_dispatch_matches_manual_grouping(): """Eval-mode grouped top-1 dispatch must equal running each row through its assigned expert individually (batch order shouldn't matter).""" B = 12 model = _routed_stage1(n_experts=4) model.eval() x_t = torch.randn(B, X_DIM) t = torch.rand(B) cond_cont, cond_cat = _cond(B) with torch.no_grad(): batched = model(x_t, t, cond_cont, cond_cat) t_emb = model.time_emb(t) c_emb = model.cond_enc(cond_cont, cond_cat) cond = torch.cat([t_emb, c_emb], dim=-1) idx = model.router.top1(cond_cont, cond_cat) manual = torch.zeros_like(x_t) for i in range(B): manual[i] = model.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0] torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4) def test_routed_denoising_mlp_predict_n_sec_shape(): B = 6 model = _routed_stage1() cond_cont, cond_cat = _cond(B) logits = model.predict_n_sec(cond_cont, cond_cat) assert logits.shape == (B, K_MAX + 1) def test_routed_denoising_mlp_has_no_pdg_embedding_weight_method(): model = _routed_stage1(pdg=5, mat=2) assert not hasattr(model, "pdg_embedding_weight") # ── RoutedSecondaryDecoder ─────────────────────────────────────────────────── def test_routed_secondary_decoder_output_shape_train_and_eval(): B = 8 decoder = _routed_sec_decoder() x_t = torch.randn(B, SEC_DIM) t = torch.rand(B) cond_cont, cond_cat = _cond(B) stage1_out = torch.randn(B, X_DIM) decoder.train() out_train = decoder(x_t, t, cond_cont, cond_cat, stage1_out) assert out_train.shape == (B, SEC_DIM) decoder.eval() with torch.no_grad(): out_eval = decoder(x_t, t, cond_cont, cond_cat, stage1_out) assert out_eval.shape == (B, SEC_DIM) def test_routed_secondary_decoder_gradients_flow(): B = 4 decoder = _routed_sec_decoder(n_experts=3) x_t = torch.randn(B, SEC_DIM) t = torch.rand(B) cond_cont, cond_cat = _cond(B) stage1_out = torch.randn(B, X_DIM) decoder.train() decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward() for name, p in decoder.named_parameters(): assert p.grad is not None, f"no grad for {name}" # ── build_models dispatch ──────────────────────────────────────────────────── def test_build_models_monolith_when_router_absent(): model_config = dict( pdg_vocab=4, mat_vocab=2, hidden_dim=32, n_blocks=2, emb_dim=16, dropout=0.1, k_max=K_MAX, ) stage1, sec_decoder = build_models(model_config) assert isinstance(stage1, DenoisingMLP) assert isinstance(sec_decoder, SecondaryDecoder) def test_build_models_monolith_when_router_disabled(): model_config = dict( pdg_vocab=4, mat_vocab=2, hidden_dim=32, n_blocks=2, emb_dim=16, dropout=0.1, k_max=K_MAX, router={"enabled": False, "type": "energy", "n_experts": 4}, ) stage1, sec_decoder = build_models(model_config) assert isinstance(stage1, DenoisingMLP) assert isinstance(sec_decoder, SecondaryDecoder) def test_build_models_routed_when_enabled(): model_config = dict( pdg_vocab=4, mat_vocab=2, emb_dim=16, dropout=0.1, k_max=K_MAX, expert_hidden_dim=16, expert_n_blocks=2, router={ "enabled": True, "type": "energy", "n_experts": 4, "temperature": 0.5, "learn_centers": True, "lambda_balance": 0.0, }, ) stage1, sec_decoder = build_models(model_config) assert isinstance(stage1, RoutedDenoisingMLP) assert isinstance(sec_decoder, RoutedSecondaryDecoder) assert len(stage1.experts) == 4 assert len(sec_decoder.experts) == 4 def test_build_models_routed_pair_is_drop_in_for_sample_flow(): """Exercise the exact calling convention giant/sample.py uses.""" from giant.sample import sample_flow, sample_secondaries model_config = dict( pdg_vocab=3, mat_vocab=2, emb_dim=16, dropout=0.1, k_max=K_MAX, expert_hidden_dim=8, expert_n_blocks=1, router={"enabled": True, "type": "energy", "n_experts": 2}, ) stage1, sec_decoder = build_models(model_config) B = 5 cond_cont, cond_cat = _cond(B, pdg=3, mat=2) stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2) assert stage1_norm.shape == (B, X_DIM) assert n_sec_pred.shape == (B,) sec_cont, sec_type_emb, sec_valid = sample_secondaries( sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2 ) assert sec_cont.shape == (B, K_MAX, 4) assert sec_valid.shape == (B, K_MAX)