"""Tests for the mixture-of-experts routing prototype (giant/model/network.py).""" import pytest import torch from giant.config import ConditioningAxisConfig from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( BLOCK_REGISTRY, TRUNK_REGISTRY, AdaLNResBlock, ComposedRouter, EnergyRouter, ExpertTrunk, FilmResBlock, PdgRouter, ProcessRouter, ROUTER_REGISTRY, ResBlock, RoutedTrunk, Stage1Model, Stage2OneShot, build_block, build_composed_router, build_expert_body, build_models, build_router, ) PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1) 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 Stage1Model( pdg_vocab=pdg, mat_vocab=mat, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG, hidden_dim=16, n_res_blocks=2, router=router, n_sec_head_k_max=K_MAX, ) def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs): router = build_router("energy", n_experts, **router_kwargs) return Stage2OneShot( pdg_vocab=pdg, mat_vocab=mat, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG, hidden_dim=16, n_res_blocks=2, generator="flow", time_dim=16, router=router, ) # ── 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 router config 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") # ── TRUNK_REGISTRY / build_expert_body ────────────────────────────────────── def test_trunk_registry_has_resmlp(): assert "resmlp" in TRUNK_REGISTRY assert TRUNK_REGISTRY["resmlp"] is ExpertTrunk def test_build_expert_body_unknown_type_raises(): try: build_expert_body("nonexistent", in_dim=4, out_dim=4, hidden_dim=8, n_blocks=1, cond_dim=4) except ValueError: return raise AssertionError("expected ValueError for unknown trunk type") # ── BLOCK_REGISTRY / build_block (gitea #34) ──────────────────────────────── def test_block_registry_has_add_film_adaln(): assert BLOCK_REGISTRY["add"] is ResBlock assert BLOCK_REGISTRY["film"] is FilmResBlock assert BLOCK_REGISTRY["adaln"] is AdaLNResBlock def test_build_block_unknown_type_raises(): try: build_block("nonexistent", dim=8, cond_dim=4) except ValueError: return raise AssertionError("expected ValueError for unknown block conditioning type") @pytest.mark.parametrize("block_type", ["add", "film", "adaln"]) def test_block_forward_shape(block_type): block = build_block(block_type, dim=8, cond_dim=4) x = torch.randn(5, 8) cond = torch.randn(5, 4) out = block(x, cond) assert out.shape == (5, 8) def test_film_res_block_output_invariant_to_cond_at_init(): """Zero-initialized film_proj means gamma=beta=0 at construction, so the output must not depend on which cond is passed in.""" block = FilmResBlock(dim=8, cond_dim=4) x = torch.randn(5, 8) cond_a = torch.randn(5, 4) cond_b = torch.randn(5, 4) torch.testing.assert_close(block(x, cond_a), block(x, cond_b)) def test_adaln_res_block_is_identity_at_init(): """Zero-initialized adaln_proj means scale=shift=gate=0 at construction, so the block must be the exact identity function (the 'Zero' in AdaLN-Zero).""" block = AdaLNResBlock(dim=8, cond_dim=4) x = torch.randn(5, 8) cond = torch.randn(5, 4) torch.testing.assert_close(block(x, cond), x) # ── EnergyRouter learn_width / learn_temperature ──────────────────────────── def test_energy_router_learn_width_matches_fixed_temperature_at_init(): """Enabling learn_width should be a no-op at init — the warm-started per-expert width must reproduce the fixed-temperature gate exactly.""" centers_init = [-1.0, 0.0, 0.5, 1.5] fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init) learned = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True) cond_cont, cond_cat = _cond(16) torch.testing.assert_close( learned.gate(cond_cont, cond_cat), fixed.gate(cond_cont, cond_cat), atol=1e-5, rtol=0, ) def test_energy_router_learn_temperature_matches_fixed_temperature_at_init(): centers_init = [-1.0, 0.0, 0.5, 1.5] fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init) learned = EnergyRouter( n_experts=4, temperature=0.3, centers_init=centers_init, learn_temperature=True, ) cond_cont, cond_cat = _cond(16) torch.testing.assert_close( learned.gate(cond_cont, cond_cat), fixed.gate(cond_cont, cond_cat), atol=1e-5, rtol=0, ) def test_energy_router_learn_width_and_temperature_mutually_exclusive_raises(): try: EnergyRouter(n_experts=4, learn_width=True, learn_temperature=True) except ValueError: return raise AssertionError("expected ValueError for learn_width and learn_temperature both set") def test_energy_router_width_ratio_bounds_must_bracket_one_raises(): try: EnergyRouter(n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0) except ValueError: return raise AssertionError("expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0") def test_energy_router_effective_width_stays_within_bounds(): router = EnergyRouter( n_experts=4, temperature=0.5, learn_width=True, width_min_ratio=0.1, width_max_ratio=10.0, ) lo, hi = 0.1 * 0.5, 10.0 * 0.5 with torch.no_grad(): router.raw_width.fill_(1e6) width = router.effective_width() assert isinstance(width, torch.Tensor) assert torch.all(width <= hi + 1e-4) with torch.no_grad(): router.raw_width.fill_(-1e6) width = router.effective_width() assert isinstance(width, torch.Tensor) assert torch.all(width >= lo - 1e-4) def test_energy_router_learn_width_gate_still_partition_of_unity(): router = EnergyRouter(n_experts=4, learn_width=True) with torch.no_grad(): router.raw_width.copy_(torch.randn(4) * 3) cond_cont, cond_cat = _cond(16) g = router.gate(cond_cont, cond_cat) torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) def test_energy_router_learn_width_hardens_when_pushed_to_floor(): """Pushing every expert's width toward the (tiny) floor should harden the gate to a one-hot at the nearest center, generalizing the fixed- temperature->0 hardening test to the per-expert path.""" router = EnergyRouter(n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0) with torch.no_grad(): router.raw_width.fill_(-1e6) cond_cont, cond_cat = _cond(16) g = router.gate(cond_cont, cond_cat) e = cond_cont[:, router.energy_idx].unsqueeze(-1) d2 = (e - router.centers.unsqueeze(0)) ** 2 onehot = torch.nn.functional.one_hot(d2.argmin(dim=-1), num_classes=4).float() torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0) def test_energy_router_own_width_controls_own_coverage_independent_of_others(): """Widening one expert's width should monotonically grow only that expert's own gate share, without needing to touch any other expert's width — the "each expert learns its own coverage independently" property this feature is meant to add.""" router = EnergyRouter(n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0]) cond_cont, cond_cat = _cond(4) cond_cont[:, 3] = 3.0 # fixed energy, unequal distance to each center shares = [] with torch.no_grad(): for raw in torch.linspace(-8.0, 8.0, 9): router.raw_width[0] = raw shares.append(router.gate(cond_cont, cond_cat)[0, 0].item()) assert all(a <= b + 1e-6 for a, b in zip(shares, shares[1:])) def test_build_router_threads_learn_width_kwargs_through(): router = build_router("energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0) assert isinstance(router, EnergyRouter) assert router.learn_width is True assert isinstance(router.raw_width, torch.nn.Parameter) assert router.raw_width.shape == (4,) def test_router_entropy_loss_is_nonnegative_bounded_scalar(): router = EnergyRouter(n_experts=4) cond_cont, cond_cat = _cond(16) loss = router.entropy_loss(cond_cont, cond_cat) assert loss.shape == () assert 0.0 <= loss.item() <= 1.0 # ── Router.combine_weights (straight-through Gumbel-softmax) ─────────────── def test_combine_weights_defaults_to_gate(): """gumbel=False (the default) must be a pure pass-through to gate().""" router = EnergyRouter(n_experts=4) cond_cont, cond_cat = _cond(16) torch.testing.assert_close( router.combine_weights(cond_cont, cond_cat), router.gate(cond_cont, cond_cat), ) def test_combine_weights_gumbel_train_mode_is_hard_one_hot(): router = EnergyRouter(n_experts=4) router.gumbel = True router.gumbel_tau = 0.5 router.train() cond_cont, cond_cat = _cond(16) weights = router.combine_weights(cond_cont, cond_cat) assert weights.shape == (16, 4) torch.testing.assert_close(weights.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) assert torch.all((weights.max(dim=-1).values - 1.0).abs() < 1e-5) def test_combine_weights_gumbel_eval_mode_falls_back_to_gate(): """No Gumbel noise at eval — combine_weights must match gate() exactly, same as the gumbel=False path, once the router is in eval mode.""" router = EnergyRouter(n_experts=4) router.gumbel = True router.eval() cond_cont, cond_cat = _cond(16) torch.testing.assert_close( router.combine_weights(cond_cont, cond_cat), router.gate(cond_cont, cond_cat), ) def test_combine_weights_gumbel_straight_through_gradient_reaches_centers(): router = EnergyRouter(n_experts=4, learn_centers=True) router.gumbel = True router.gumbel_tau = 0.5 router.train() cond_cont, cond_cat = _cond(16) weights = router.combine_weights(cond_cont, cond_cat) weights.sum().backward() assert router.centers.grad is not None assert torch.any(router.centers.grad != 0.0) def test_build_router_from_cfg_sets_gumbel_from_config(): from giant.model.network import _build_router_from_cfg router = _build_router_from_cfg( {"enabled": True, "type": "energy", "n_experts": 4, "gumbel": True}, pdg_vocab=3, mat_vocab=2, ) assert router.gumbel is True router_off = _build_router_from_cfg( {"enabled": True, "type": "energy", "n_experts": 4}, pdg_vocab=3, mat_vocab=2, ) assert router_off.gumbel is False def test_build_router_from_cfg_sets_gumbel_for_composed_router(): from giant.model.network import _build_router_from_cfg router = _build_router_from_cfg( { "enabled": True, "type": "composed", "gumbel": True, "axis0_type": "energy", "axis0_n_experts": 4, "axis1_type": "pdg", "axis1_n_experts": 3, }, pdg_vocab=5, mat_vocab=2, ) assert isinstance(router, ComposedRouter) assert router.gumbel is True def test_routed_stage1_forward_runs_with_gumbel_enabled(): """End-to-end forward through _route_forward's train branch with straight-through Gumbel-softmax combine weights enabled.""" B = 8 model = _routed_stage1(n_experts=3) model.trunk.router.gumbel = True model.trunk.router.gumbel_tau = 0.5 model.train() x_t = torch.randn(B, X_DIM) t = torch.rand(B) cond_cont, cond_cat = _cond(B) out = model(x_t, cond_cont, cond_cat, t=t) assert out.shape == (B, X_DIM) assert torch.isfinite(out).all() # ── 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 _nested_cfg( pdg_vocab, mat_vocab, stage1_router=None, stage2_router=None, particle_type="physical", material_type="physical", **overrides, ): """Minimal new-shape (v0.3.0) model_config for build_models, with optional router sub-blocks. `overrides` deep-patches stage1_model.""" stage1_model = { "active": True, "generator": "flow", "hidden_dim": 16, "n_res_blocks": 2, "dropout": 0.0, "flow": {"time_dim": 16}, "router": stage1_router or {"enabled": False}, } stage1_model.update(overrides) return { "pdg_vocab": pdg_vocab, "mat_vocab": mat_vocab, "conditioning": { "out_dim": 32, "particle": {"type": particle_type, "emb_dim": 8, "n_layers": 1}, "material": {"type": material_type, "emb_dim": 8, "n_layers": 1}, }, "stage1_model": stage1_model, "stage2_model": { "active": True, "decoder": "one_shot", "generator": "flow", "hidden_dim": 16, "n_res_blocks": 2, "dropout": 0.0, "k_max": K_MAX, "context_dim": 16, "n_sec": {"mode": "head"}, "flow": {"time_dim": 16}, "router": stage2_router or {"enabled": False, "tie_to_stage1": False}, }, } def test_build_models_routed_with_pdg_router(): cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, particle_type="embedding", material_type="embedding", stage1_router={"enabled": True, "type": "pdg", "n_experts": 3}, ) models = build_models(cfg) stage1 = models["stage1"] assert isinstance(stage1, Stage1Model) assert isinstance(stage1.trunk, RoutedTrunk) assert isinstance(stage1.trunk.router, PdgRouter) assert len(stage1.trunk.experts) == 3 assert stage1.trunk.router.pdg_emb.num_embeddings == 4 def test_build_models_rejects_pdg_router_with_physical_conditioning(): """conditioning.particle.type="physical" is meant to generalize beyond the training PDG vocab; PdgRouter always uses a training-vocab nn.Embedding regardless of conditioning, so the combination must raise rather than silently building a model that can't actually generalize the way it claims to.""" cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, stage1_router={"enabled": True, "type": "pdg", "n_experts": 3}, ) with pytest.raises(ValueError, match="physical"): build_models(cfg) def test_build_models_rejects_composed_router_with_pdg_axis_and_physical_conditioning(): cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, stage1_router={ "enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 2, "axis1_type": "pdg", "axis1_n_experts": 3, }, ) with pytest.raises(ValueError, match="physical"): build_models(cfg) # ── 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(): cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, particle_type="embedding", material_type="embedding", stage1_router={ "enabled": True, "type": "process", "n_experts": 3, "lambda_proc": 1.0, }, ) models = build_models(cfg) stage1 = models["stage1"] assert isinstance(stage1, Stage1Model) assert isinstance(stage1.trunk, RoutedTrunk) assert isinstance(stage1.trunk.router, ProcessRouter) assert len(stage1.trunk.experts) == 3 assert stage1.trunk.router.pdg_emb.num_embeddings == 4 assert stage1.trunk.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(): cfg = _nested_cfg( pdg_vocab=5, mat_vocab=2, particle_type="embedding", material_type="embedding", stage1_router={ "enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4, "axis1_type": "pdg", "axis1_n_experts": 3, }, stage2_router={ "enabled": True, "tie_to_stage1": False, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4, "axis1_type": "pdg", "axis1_n_experts": 3, }, ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] assert stage1 is not None and stage2 is not None assert isinstance(stage1.trunk, RoutedTrunk) assert isinstance(stage1.trunk.router, ComposedRouter) assert len(stage1.trunk.experts) == 12 assert len(stage2.trunk.experts) == 12 # stage1 and stage2 must not share router weights when tie_to_stage1 is # false (same convention as v0.2's two-independent-routers behaviour). assert stage1.trunk.router is not stage2.trunk.router def test_build_models_routed_stage2_ties_to_stage1_router(): cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, stage1_router={"enabled": True, "type": "energy", "n_experts": 3}, stage2_router={ "enabled": True, "tie_to_stage1": True, "type": "energy", "n_experts": 3, }, ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] assert stage1 is not None and stage2 is not None assert stage1.trunk.router is stage2.trunk.router def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow(): from giant.sample import sample_flow, sample_secondaries cfg = _nested_cfg( pdg_vocab=3, mat_vocab=2, # PdgRouter (axis1) always builds its own training-vocab embedding, # incompatible with conditioning.particle.type="physical" (the # _nested_cfg default) — see _check_router_conditioning_compat. particle_type="embedding", material_type="embedding", stage1_router={ "enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 2, "axis1_type": "pdg", "axis1_n_experts": 2, }, ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] assert stage1 is not None and stage2 is not None 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) # A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) — # sample_flow returns n_sec_pred=None here, and n_sec must be # asked of stage2 instead, using the just-sampled stage1_norm as context. assert n_sec_pred is None n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1) assert n_sec_pred.shape == (B,) sec_cont, sec_type_emb, sec_valid = sample_secondaries( stage2, 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) # ── Stage1Model with a routed trunk ───────────────────────────────────────── def test_routed_stage1_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, cond_cont, cond_cat, t=t) assert out_train.shape == (B, X_DIM) model.eval() with torch.no_grad(): out_eval = model(x_t, cond_cont, cond_cat, t=t) assert out_eval.shape == (B, X_DIM) def test_routed_stage1_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, cond_cont, cond_cat, t=t).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_stage1_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, cond_cont, cond_cat, t=t) 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.trunk.router.top1(cond_cont, cond_cat) manual = torch.zeros_like(x_t) for i in range(B): manual[i] = model.trunk.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_stage1_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_stage1_has_no_pdg_embedding_weight_method(): model = _routed_stage1(pdg=5, mat=2) assert not hasattr(model, "pdg_embedding_weight") # ── Stage2OneShot with a routed trunk ──────────────────────────────────────── def test_routed_stage2_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, cond_cont, cond_cat, stage1_out, t=t) assert out_train.shape == (B, SEC_DIM) decoder.eval() with torch.no_grad(): out_eval = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t) assert out_eval.shape == (B, SEC_DIM) def test_routed_stage2_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() flow_loss = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t).sum() nsec_loss = decoder.predict_n_sec(cond_cont, cond_cat, stage1_out).sum() (flow_loss + nsec_loss).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(): cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] assert isinstance(stage1, Stage1Model) assert isinstance(stage2, Stage2OneShot) assert not isinstance(stage1.trunk, RoutedTrunk) assert not isinstance(stage2.trunk, RoutedTrunk) assert isinstance(stage1.trunk, ExpertTrunk) assert isinstance(stage2.trunk, ExpertTrunk) def test_build_models_monolith_when_router_disabled(): cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, stage1_router={"enabled": False, "type": "energy", "n_experts": 4}, ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] assert stage1 is not None and stage2 is not None assert not isinstance(stage1.trunk, RoutedTrunk) assert not isinstance(stage2.trunk, RoutedTrunk) assert isinstance(stage1.trunk, ExpertTrunk) assert isinstance(stage2.trunk, ExpertTrunk) def test_build_models_routed_when_enabled(): cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, stage1_router={ "enabled": True, "type": "energy", "n_experts": 4, "temperature": 0.5, "learn_centers": True, "lambda_balance": 0.0, }, stage2_router={ "enabled": True, "tie_to_stage1": False, "type": "energy", "n_experts": 4, "temperature": 0.5, "learn_centers": True, }, ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] assert stage1 is not None and stage2 is not None assert isinstance(stage1.trunk, RoutedTrunk) assert isinstance(stage2.trunk, RoutedTrunk) assert len(stage1.trunk.experts) == 4 assert len(stage2.trunk.experts) == 4 def test_build_models_explicit_resmlp_trunk_type_matches_default(): """stage1_model.trunk.type = 'resmlp' is the default's spelled-out equivalent, not a behaviour change — gitea #33.""" default_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2) explicit_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp"}) default_stage1 = build_models(default_cfg)["stage1"] explicit_stage1 = build_models(explicit_cfg)["stage1"] assert default_stage1 is not None and explicit_stage1 is not None assert type(default_stage1.trunk) is type(explicit_stage1.trunk) is ExpertTrunk assert default_stage1.trunk.input_proj.weight.shape == explicit_stage1.trunk.input_proj.weight.shape default_params = sum(p.numel() for p in default_stage1.parameters()) explicit_params = sum(p.numel() for p in explicit_stage1.parameters()) assert default_params == explicit_params def test_build_models_explicit_add_block_conditioning_matches_default(): """stage1_model.trunk.block_conditioning = 'add' is the default's spelled-out equivalent, not a behaviour change — gitea #34.""" default_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2) explicit_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp", "block_conditioning": "add"}) default_stage1 = build_models(default_cfg)["stage1"] explicit_stage1 = build_models(explicit_cfg)["stage1"] assert default_stage1 is not None and explicit_stage1 is not None assert type(default_stage1.trunk.blocks[0]) is type(explicit_stage1.trunk.blocks[0]) is ResBlock default_params = sum(p.numel() for p in default_stage1.parameters()) explicit_params = sum(p.numel() for p in explicit_stage1.parameters()) assert default_params == explicit_params @pytest.mark.parametrize("block_type,cls", [("film", FilmResBlock), ("adaln", AdaLNResBlock)]) def test_build_models_selects_block_conditioning(block_type, cls): cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp", "block_conditioning": block_type}) stage1 = build_models(cfg)["stage1"] assert stage1 is not None assert isinstance(stage1.trunk, ExpertTrunk) assert all(isinstance(b, cls) for b in stage1.trunk.blocks) def test_build_models_routed_trunk_uses_block_conditioning_for_every_expert(): cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp", "block_conditioning": "film"}, stage1_router={"enabled": True, "type": "energy", "n_experts": 3}, ) stage1 = build_models(cfg)["stage1"] assert stage1 is not None assert isinstance(stage1.trunk, RoutedTrunk) assert len(stage1.trunk.experts) == 3 for expert in stage1.trunk.experts: assert all(isinstance(b, FilmResBlock) for b in expert.blocks) 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 cfg = _nested_cfg( pdg_vocab=3, mat_vocab=2, stage1_router={"enabled": True, "type": "energy", "n_experts": 2}, ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] assert stage1 is not None and stage2 is not None 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) # A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) — # sample_flow returns n_sec_pred=None here, and n_sec must be # asked of stage2 instead, using the just-sampled stage1_norm as context. assert n_sec_pred is None n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1) assert n_sec_pred.shape == (B,) sec_cont, sec_type_emb, sec_valid = sample_secondaries( stage2, 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)