configs/router_energy_n10_gumbel_centers_physical.toml

@@ -0,0 +1,22 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning = "physical"
+dropout      = 0.0
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = true
+gumbel            = true

configs/router_energy_n10_gumbel_centers_temp_physical.toml

@@ -0,0 +1,23 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning = "physical"
+dropout      = 0.0
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = true
+learn_temperature = true
+gumbel            = true

configs/router_energy_n10_gumbel_physical.toml

@@ -0,0 +1,22 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning = "physical"
+dropout      = 0.0
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = false
+gumbel            = true

giant/config.py

@@ -99,6 +99,19 @@ DEFAULT_CONFIG: dict = {
             # failure mode. Off by default; bounding above is the primary
             # defense. See giant.model.network.Router.entropy_loss.
             "lambda_entropy": 0.0,
+            # Opt-in straight-through Gumbel-softmax train-time combine weights
+            # (see giant.model.network.Router.combine_weights): the training
+            # forward pass samples a hard one-hot combination — matching
+            # eval-time top-1 dispatch exactly — while the backward pass still
+            # flows a smooth gradient to every expert. Targets the train/eval
+            # mismatch identified as a likely contributor to experts
+            # overlapping instead of partitioning (see CLAUDE.md roadmap).
+            # gumbel_tau_start/_end are annealed linearly over training
+            # (giant.train._gumbel_tau); off by default, no effect unless
+            # gumbel = true.
+            "gumbel": False,
+            "gumbel_tau_start": 1.0,
+            "gumbel_tau_end": 0.1,
             "emb_dim": 8,  # process/pdg-router kwarg: own pdg(/mat) embedding width
             "hidden_dim": 64,  # process-router kwarg: its classifier's hidden width
             "lambda_proc": 0.0,  # process-router kwarg: supervised process-CE weight
@@ -326,6 +339,29 @@ def _router_candidate(train, model):
     return f"r-{router['type']}{router['n_experts']}"
 
 
+def _router_flag_candidate(field, token_map):
+    """Candidate factory for a boolean `model.router` sub-field.
+
+    Gated on `router.enabled` like `_router_candidate` (a disabled router's
+    sub-fields are meaningless), then omitted unless `field` differs from
+    its DEFAULT_CONFIG value — same "only show non-default" rule as every
+    other candidate. `token_map` need only cover the non-default value(s),
+    since the default value always yields None.
+    """
+
+    def _candidate(train, model):
+        router = model["router"]
+        default_router = DEFAULT_CONFIG["model"]["router"]
+        if router["enabled"] == default_router["enabled"]:
+            return None
+        value = router[field]
+        if value == default_router[field]:
+            return None
+        return token_map[value]
+
+    return _candidate
+
+
 def _conditioning_candidate(train, model):
     if model["conditioning"] == DEFAULT_CONFIG["model"]["conditioning"]:
         return None
@@ -347,6 +383,10 @@ def _default_field_candidate(section_key, field, prefix):
 _OUT_DIR_NAME_CANDIDATES = [
     ("mode", _mode_candidate),
     ("router", _router_candidate),
+    ("gumbel", _router_flag_candidate("gumbel", {True: "gum"})),
+    ("learn_centers", _router_flag_candidate("learn_centers", {False: "nolc"})),
+    ("learn_width", _router_flag_candidate("learn_width", {True: "lw"})),
+    ("learn_temperature", _router_flag_candidate("learn_temperature", {True: "lt"})),
     ("conditioning", _conditioning_candidate),
     ("hidden_dim", _default_field_candidate("model", "hidden_dim", "h")),
     ("n_blocks", _default_field_candidate("model", "n_blocks", "b")),

giant/model/network.py

@@ -537,11 +537,49 @@ class Router(nn.Module):
     def __init__(self, n_experts: int) -> None:
         super().__init__()
         self.n_experts = n_experts
+        # Opt-in straight-through Gumbel-softmax combine weights (see
+        # combine_weights below) — off by default, set from model.router.gumbel
+        # by _build_router_from_cfg. gumbel_tau is annealed per training step
+        # by giant.train (model.router.gumbel_tau_start/_end); neither is an
+        # nn.Parameter/buffer since neither is learned or needs checkpointing —
+        # the tau schedule is deterministic in global_step, so it recomputes
+        # correctly on resume.
+        self.gumbel = False
+        self.gumbel_tau = 1.0
 
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         """(B, n_experts) soft weights, rows summing to 1."""
         raise NotImplementedError
 
+    def combine_weights(
+        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
+    ) -> torch.Tensor:
+        """(B, n_experts) train-time expert-combination weights.
+
+        Default (`gumbel=False`): identical to `gate()` — the original dense
+        soft-mixture combination. Opt-in straight-through Gumbel-softmax
+        (`gumbel=True`, train mode only): samples a Gumbel-perturbed
+        categorical draw from the same distribution `gate()` defines
+        (`log(gate())` is a valid unnormalized-logit input to
+        `F.gumbel_softmax` since softmax is shift-invariant, so no subclass
+        needs to expose separate pre-softmax logits), then hardens it to a
+        one-hot vector on the forward pass while keeping the soft sample's
+        gradient on the backward pass. This makes the training-time forward
+        combination match eval-time top-1 dispatch exactly (one expert's
+        output, unweighted) instead of the smooth blend `gate()` gives —
+        intended to close the train/eval mismatch identified as a likely
+        cause of experts overlapping instead of partitioning (see the
+        router_gating write-up referenced in CLAUDE.md's roadmap).
+        `gate()` itself is untouched and still backs `balance_loss`/
+        `entropy_loss`/`gate_stats`, so those diagnostics keep reading the
+        smooth distribution rather than a noisy sample.
+        """
+        probs = self.gate(cond_cont, cond_cat)
+        if not (self.gumbel and self.training):
+            return probs
+        log_probs = torch.log(probs.clamp_min(1e-8))
+        return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1)
+
     def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         """(B,) hard expert index, used for eval-time grouped dispatch."""
         return self.gate(cond_cont, cond_cat).argmax(dim=-1)
@@ -970,13 +1008,17 @@ def _route_forward(
 ) -> torch.Tensor:
     """Shared dispatch for both Routed* trunks.
 
-    Train mode: full soft mixture `sum_i gate_i * expert_i(x)` — fully
-    differentiable, N-expert compute. Eval mode: grouped top-1 dispatch —
-    each row runs exactly one (small) expert, which is the actual source
-    of the per-call speedup this architecture is for.
+    Train mode: full mixture `sum_i weight_i * expert_i(x)` — always
+    N-expert dense compute, fully differentiable. `weight` is
+    `router.combine_weights(...)`: the plain soft `gate()` by default, or (see
+    `Router.combine_weights`) a straight-through Gumbel-softmax one-hot sample
+    when `router.gumbel` is enabled — either way, no change to the compute
+    cost of this branch. Eval mode: grouped top-1 dispatch — each row runs
+    exactly one (small) expert, which is the actual source of the per-call
+    speedup this architecture is for.
     """
     if training:
-        weights = router.gate(cond_cont, cond_cat)  # (B, n_experts)
+        weights = router.combine_weights(cond_cont, cond_cat)  # (B, n_experts)
         out = torch.zeros_like(x)
         for i, expert in enumerate(experts):
             out = out + weights[:, i : i + 1] * expert(x, cond)
@@ -1191,10 +1233,19 @@ def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) ->
 
     `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys
     (see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair.
+
+    `gumbel` is set as a post-construction attribute here rather than a
+    per-subclass constructor kwarg, same reasoning as `lambda_balance`/
+    `lambda_proc`/`lambda_entropy` living in `router_cfg` without being a
+    `Router` subclass constructor param: it's a training-time toggle shared by
+    every router type, not a per-type hyperparameter (`build_router`'s
+    kwarg-filtering would otherwise just silently drop it).
     """
     shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
     if router_cfg["type"] == "composed":
-        return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab)
+        router = build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab)
+        router.gumbel = bool(router_cfg.get("gumbel", False))
+        return router
     router_kwargs = {
         k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
     }
@@ -1204,7 +1255,9 @@ def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) ->
     # vocab, same as the trunk's ConditionEncoder.
     router_kwargs.setdefault("pdg_vocab", pdg_vocab)
     router_kwargs.setdefault("mat_vocab", mat_vocab)
-    return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
+    router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
+    router.gumbel = bool(router_cfg.get("gumbel", False))
+    return router
 
 
 def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:

giant/pipeline.py

@@ -417,6 +417,8 @@ def run_train_job(
         lambda_balance=router_cfg.get("lambda_balance", 0.0),
         lambda_proc=router_cfg.get("lambda_proc", 0.0),
         lambda_entropy=router_cfg.get("lambda_entropy", 0.0),
+        gumbel_tau_start=router_cfg.get("gumbel_tau_start", 1.0),
+        gumbel_tau_end=router_cfg.get("gumbel_tau_end", 0.1),
         normalizer_dict={
             "cond": cond_norm.to_dict(),
             "target": tgt_norm.to_dict(),

giant/train.py

@@ -114,6 +114,85 @@ def _update_ema(
         ema_p.mul_(decay).add_(p, alpha=1 - decay)
 
 
+def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -> float:
+    """Linear anneal of the straight-through Gumbel-softmax temperature.
+
+    Deterministic in `step`/`total_steps` alone (no extra state), so it
+    recomputes correctly on `--resume` from a checkpoint's saved `global_step`
+    without needing to persist anything new (see
+    giant.model.network.Router.combine_weights).
+    """
+    progress = min(step / max(total_steps, 1), 1.0)
+    return tau_start + (tau_end - tau_start) * progress
+
+
+def _wandb_run_config(
+    *,
+    mode: str,
+    epochs: int,
+    lr: float,
+    warmup_epochs: int,
+    weight_decay: float,
+    ema_decay: float,
+    lambda_nsec: float,
+    lambda_s2: float,
+    lambda_balance: float,
+    lambda_proc: float,
+    lambda_entropy: float,
+    gumbel_tau_start: float,
+    gumbel_tau_end: float,
+    n_critic: int,
+    gp_weight: float,
+    model_config: dict | None,
+    stage1_params: int,
+    sec_decoder_params: int,
+    critic_params: int,
+    sec_critic_params: int,
+    total_params: int,
+) -> dict:
+    """Build the dict logged as a wandb run's `config`.
+
+    Router-only knobs (`lambda_balance`/`lambda_proc`/`lambda_entropy`/
+    `gumbel_tau_start`/`gumbel_tau_end`) and WGAN-only knobs (`n_critic`/
+    `gp_weight`) are omitted unless actually active, so a run's wandb config
+    doesn't imply hyperparameters from an inactive code path (a disabled
+    router's fine-tuning knobs, or GAN critic settings for a flow/DDPM run).
+    The full `model_config` (including its `router` sub-dict, whatever the
+    router type/state) is always included, so no information is lost — this
+    only trims the flattened top-level convenience duplicates.
+    """
+    router_enabled = bool((model_config or {}).get("router", {}).get("enabled", False))
+    cfg = {
+        "mode": mode,
+        "epochs": epochs,
+        "lr": lr,
+        "warmup_epochs": warmup_epochs,
+        "weight_decay": weight_decay,
+        "ema_decay": ema_decay,
+        "lambda_nsec": lambda_nsec,
+        "lambda_s2": lambda_s2,
+        "model": model_config or {},
+        "stage1_params": stage1_params,
+        "sec_decoder_params": sec_decoder_params,
+        "critic_params": critic_params,
+        "sec_critic_params": sec_critic_params,
+        "total_params": total_params,
+    }
+    if router_enabled:
+        cfg.update(
+            {
+                "lambda_balance": lambda_balance,
+                "lambda_proc": lambda_proc,
+                "lambda_entropy": lambda_entropy,
+                "gumbel_tau_start": gumbel_tau_start,
+                "gumbel_tau_end": gumbel_tau_end,
+            }
+        )
+    if mode == "wgan":
+        cfg.update({"n_critic": n_critic, "gp_weight": gp_weight})
+    return cfg
+
+
 def _compute_losses(
     stage1_model: torch.nn.Module,
     sec_decoder: torch.nn.Module,
@@ -347,6 +426,8 @@ def train(
     lambda_balance: float = 0.0,
     lambda_proc: float = 0.0,
     lambda_entropy: float = 0.0,
+    gumbel_tau_start: float = 1.0,
+    gumbel_tau_end: float = 0.1,
     normalizer_dict: dict | None = None,
     pdg_map: dict | None = None,
     mat_map: dict | None = None,
@@ -398,27 +479,29 @@ def train(
             name=wandb_run_name or out_dir.name,
             id=out_dir.name,
             resume="allow",
-            config={
-                "mode": mode,
-                "epochs": epochs,
-                "lr": lr,
-                "warmup_epochs": warmup_epochs,
-                "weight_decay": weight_decay,
-                "ema_decay": ema_decay,
-                "lambda_nsec": lambda_nsec,
-                "lambda_s2": lambda_s2,
-                "lambda_balance": lambda_balance,
-                "lambda_proc": lambda_proc,
-                "lambda_entropy": lambda_entropy,
-                "n_critic": n_critic,
-                "gp_weight": gp_weight,
-                "model": model_config or {},
-                "stage1_params": stage1_params,
-                "sec_decoder_params": sec_decoder_params,
-                "critic_params": critic_params,
-                "sec_critic_params": sec_critic_params,
-                "total_params": total_params,
-            },
+            config=_wandb_run_config(
+                mode=mode,
+                epochs=epochs,
+                lr=lr,
+                warmup_epochs=warmup_epochs,
+                weight_decay=weight_decay,
+                ema_decay=ema_decay,
+                lambda_nsec=lambda_nsec,
+                lambda_s2=lambda_s2,
+                lambda_balance=lambda_balance,
+                lambda_proc=lambda_proc,
+                lambda_entropy=lambda_entropy,
+                gumbel_tau_start=gumbel_tau_start,
+                gumbel_tau_end=gumbel_tau_end,
+                n_critic=n_critic,
+                gp_weight=gp_weight,
+                model_config=model_config,
+                stage1_params=stage1_params,
+                sec_decoder_params=sec_decoder_params,
+                critic_params=critic_params,
+                sec_critic_params=sec_critic_params,
+                total_params=total_params,
+            ),
         )
 
     stage1_model = stage1_model.to(device)
@@ -596,6 +679,13 @@ def train(
                 dynamic_ncols=True,
             )
             for batch in bar:
+                if has_router:
+                    gumbel_tau = _gumbel_tau(
+                        global_step, total_steps, gumbel_tau_start, gumbel_tau_end
+                    )
+                    stage1_model.router.gumbel_tau = gumbel_tau
+                    sec_decoder.router.gumbel_tau = gumbel_tau
+
                 if mode == "wgan":
                     assert (
                         critic is not None
@@ -737,6 +827,7 @@ def train(
                             )
                         log_payload["batch/router_s1_entropy"] = s1_entropy.item()
                         log_payload["batch/router_s2_entropy"] = s2_entropy.item()
+                        log_payload["batch/gumbel_tau"] = gumbel_tau
                     wandb_run.log(log_payload, step=global_step)
 
                 if shutdown.requested:

tests/test_config.py

@@ -137,6 +137,16 @@ def test_resolve_expert_dims_missing_keys_also_inherit():
     assert (hidden_dim, n_blocks) == (512, 6)
 
 
+def test_default_config_gumbel_router_defaults_off():
+    # Straight-through Gumbel-softmax combine weights (giant.model.network.
+    # Router.combine_weights) must be opt-in — existing routed configs and
+    # checkpoints should be unaffected unless gumbel is explicitly enabled.
+    router_cfg = gconfig.DEFAULT_CONFIG["model"]["router"]
+    assert router_cfg["gumbel"] is False
+    assert router_cfg["gumbel_tau_start"] == 1.0
+    assert router_cfg["gumbel_tau_end"] == 0.1
+
+
 def test_resolve_expert_dims_explicit_override_wins():
     router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 3}
     hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
@@ -197,6 +207,64 @@ def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefau
     assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
 
 
+def test_default_out_dir_name_router_gumbel_shown_when_enabled():
+    cfg = _default_cfg(
+        router={"enabled": True, "type": "energy", "n_experts": 8, "gumbel": True}
+    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_gum"
+
+
+def test_default_out_dir_name_router_gumbel_omitted_when_router_disabled():
+    cfg = _default_cfg(
+        router={"enabled": False, "type": "energy", "n_experts": 8, "gumbel": True}
+    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
+
+
+def test_default_out_dir_name_router_learn_centers_shown_only_when_disabled():
+    cfg_default = _default_cfg(
+        router={"enabled": True, "type": "energy", "n_experts": 8}
+    )
+    assert (
+        gconfig.default_out_dir_name(cfg_default, now=_NOW) == "20260729_1430_r-energy8"
+    )
+
+    cfg_off = _default_cfg(
+        router={
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 8,
+            "learn_centers": False,
+        }
+    )
+    assert (
+        gconfig.default_out_dir_name(cfg_off, now=_NOW)
+        == "20260729_1430_r-energy8_nolc"
+    )
+
+
+def test_default_out_dir_name_router_learn_width_and_temperature_shown():
+    cfg = _default_cfg(
+        router={
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 8,
+            "learn_width": True,
+        }
+    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_lw"
+
+    cfg2 = _default_cfg(
+        router={
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 8,
+            "learn_temperature": True,
+        }
+    )
+    assert gconfig.default_out_dir_name(cfg2, now=_NOW) == "20260729_1430_r-energy8_lt"
+
+
 def test_default_out_dir_name_mode_shown_bare_no_prefix():
     cfg = _default_cfg(mode="wgan")
     assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan"

tests/test_router.py

@@ -286,6 +286,110 @@ def test_router_entropy_loss_is_nonnegative_bounded_scalar():
     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_denoising_mlp_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.router.gumbel = True
+    model.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, t, cond_cont, cond_cat)
+    assert out.shape == (B, X_DIM)
+    assert torch.isfinite(out).all()
+
+
 # ── PdgRouter ────────────────────────────────────────────────────────────────
 
 

tests/test_train.py

@@ -0,0 +1,97 @@
+"""Tests for giant/train.py helpers."""
+
+from giant.train import _gumbel_tau, _wandb_run_config
+
+
+def test_gumbel_tau_at_step_zero_is_start():
+    assert _gumbel_tau(0, 1000, 1.0, 0.1) == 1.0
+
+
+def test_gumbel_tau_at_total_steps_is_end():
+    assert abs(_gumbel_tau(1000, 1000, 1.0, 0.1) - 0.1) < 1e-9
+
+
+def test_gumbel_tau_interpolates_linearly_midway():
+    assert abs(_gumbel_tau(500, 1000, 1.0, 0.1) - 0.55) < 1e-9
+
+
+def test_gumbel_tau_clamps_beyond_total_steps():
+    assert _gumbel_tau(5000, 1000, 1.0, 0.1) == _gumbel_tau(1000, 1000, 1.0, 0.1)
+
+
+def test_gumbel_tau_handles_zero_total_steps():
+    # total_steps=0 is guarded to 1 internally: step=0 gives zero progress
+    # (still tau_start), any step>=1 immediately clamps to full progress.
+    assert _gumbel_tau(0, 0, 1.0, 0.1) == 1.0
+    assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9
+
+
+def _base_wandb_kwargs(**overrides):
+    kwargs = dict(
+        mode="flow",
+        epochs=30,
+        lr=3e-4,
+        warmup_epochs=3,
+        weight_decay=0.01,
+        ema_decay=0.9999,
+        lambda_nsec=0.1,
+        lambda_s2=1.0,
+        lambda_balance=0.035,
+        lambda_proc=0.0,
+        lambda_entropy=0.0,
+        gumbel_tau_start=1.0,
+        gumbel_tau_end=0.1,
+        n_critic=5,
+        gp_weight=10.0,
+        model_config={"router": {"enabled": False}},
+        stage1_params=100,
+        sec_decoder_params=50,
+        critic_params=0,
+        sec_critic_params=0,
+        total_params=150,
+    )
+    kwargs.update(overrides)
+    return kwargs
+
+
+def test_wandb_run_config_omits_router_knobs_when_router_disabled():
+    cfg = _wandb_run_config(**_base_wandb_kwargs())
+    for key in (
+        "lambda_balance",
+        "lambda_proc",
+        "lambda_entropy",
+        "gumbel_tau_start",
+        "gumbel_tau_end",
+    ):
+        assert key not in cfg
+    # still present, nested, regardless of router state
+    assert cfg["model"] == {"router": {"enabled": False}}
+
+
+def test_wandb_run_config_includes_router_knobs_when_router_enabled():
+    cfg = _wandb_run_config(
+        **_base_wandb_kwargs(model_config={"router": {"enabled": True}})
+    )
+    assert cfg["lambda_balance"] == 0.035
+    assert cfg["lambda_proc"] == 0.0
+    assert cfg["lambda_entropy"] == 0.0
+    assert cfg["gumbel_tau_start"] == 1.0
+    assert cfg["gumbel_tau_end"] == 0.1
+
+
+def test_wandb_run_config_omits_wgan_knobs_when_mode_is_not_wgan():
+    cfg = _wandb_run_config(**_base_wandb_kwargs(mode="flow"))
+    assert "n_critic" not in cfg
+    assert "gp_weight" not in cfg
+
+
+def test_wandb_run_config_includes_wgan_knobs_when_mode_is_wgan():
+    cfg = _wandb_run_config(**_base_wandb_kwargs(mode="wgan"))
+    assert cfg["n_critic"] == 5
+    assert cfg["gp_weight"] == 10.0
+
+
+def test_wandb_run_config_handles_missing_model_config():
+    cfg = _wandb_run_config(**_base_wandb_kwargs(model_config=None))
+    assert cfg["model"] == {}
+    assert "lambda_balance" not in cfg