Add learnable per-expert width and shared temperature to EnergyRouter

EnergyRouter's gate sharpness was a single fixed temperature shared by
every expert, with no way for an expert to independently learn how much
of the energy axis it covers. Adds two mutually exclusive, default-off
modes: learn_width (per-expert learnable width) and learn_temperature
(single learnable shared scalar), both bounded via a sigmoid
interpolation warm-started to reproduce today's fixed-temperature gate
exactly at init, to compare against each other without risking the
unbounded-width collapse failure mode. Also promotes gate_stats's
entropy into a generic, optional Router.entropy_loss (lambda_entropy) as
a secondary guard against all experts' widths co-inflating together.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

tests/test_router.py

@@ -141,6 +141,151 @@ def test_build_router_unknown_type_raises():
     raise AssertionError("expected ValueError for unknown router type")
 
 
+# ── 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
+
+
 # ── PdgRouter ────────────────────────────────────────────────────────────────