Merge pull request 'Add "none" variants for router, history, and trunk (gitea #45)' (#73) from fix/issue-45 into master

Reviewed-on: #73

giant/model/history.py

@@ -63,6 +63,23 @@ def build_history(name: str, in_dim: int, out_dim: int, **kwargs) -> HistoryEnco
     return cls(in_dim, out_dim, **filtered)
 
 
+@register_history("none")
+class NoHistory(HistoryEncoder):
+    """No history signal at all — ignores feat/has_prev entirely and always
+    returns zeros. Ablates whether the AR decoder's history conditioning is
+    earning its parameters. `init_cache`/`step` use the base class's O(1)
+    defaults unmodified (this encoder's own `forward` is already O(1) per
+    call regardless of prefix length)."""
+
+    def __init__(self, in_dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.out_dim = out_dim
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        B, K, _ = feat.shape
+        return torch.zeros(B, K, self.out_dim, device=feat.device, dtype=feat.dtype)
+
+
 @register_history("markov")
 class MarkovHistory(HistoryEncoder):
     """Summarizes the previous secondary's own `(energy_fraction, direction,

giant/model/network.py

@@ -15,6 +15,7 @@ from giant.model.history import (
     AttentionHistory,
     HistoryEncoder,
     MarkovHistory,
+    NoHistory,
     _CausalAttnBlock,
     build_history,
     register_history,
@@ -54,6 +55,7 @@ from giant.model.routers import (
     ROUTER_REGISTRY,
     ComposedRouter,
     EnergyRouter,
+    NoneRouter,
     PdgRouter,
     ProcessRouter,
     Router,
@@ -67,6 +69,7 @@ from giant.model.routers import (
 from giant.model.trunks import (
     TRUNK_REGISTRY,
     ExpertTrunk,
+    LinearTrunk,
     RoutedTrunk,
     Trunk,
     _route_forward,
@@ -90,7 +93,10 @@ __all__ = [
     "FlowObjective",
     "HISTORY_REGISTRY",
     "HistoryEncoder",
+    "LinearTrunk",
     "MarkovHistory",
+    "NoHistory",
+    "NoneRouter",
     "OBJECTIVE_REGISTRY",
     "Objective",
     "PdgRouter",

giant/model/routers.py

@@ -144,6 +144,24 @@ def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
     return math.log(p / (1 - p))
 
 
+@register_router("none")
+class NoneRouter(Router):
+    """Uniform 1/n_experts gate — no learned routing signal at all.
+
+    Still builds n_experts expert trunks via RoutedTrunk (same parameter
+    budget as a real router), but every row gets an identical weight
+    regardless of conditioning. Ablates whether the *learned routing
+    signal* — as opposed to simply having multiple experts — is earning
+    its parameters. `top1()` (the base class default) always dispatches to
+    expert 0 (argmax of a uniform vector), which still exercises
+    RoutedTrunk's real per-expert grouped-dispatch code path at eval time.
+    """
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        B = cond_cont.shape[0]
+        return torch.full((B, self.n_experts), 1.0 / self.n_experts, device=cond_cont.device)
+
+
 @register_router("energy")
 class EnergyRouter(Router):
     """Soft turn-on gate over normalized pre-step log-energy.

giant/model/trunks.py

@@ -94,6 +94,49 @@ class ExpertTrunk(nn.Module):
         return self.out_proj(x)
 
 
+@register_trunk("linear")
+class LinearTrunk(nn.Module):
+    """`nn.Linear(in_dim + cond_dim, out_dim)` over `concat([x, cond])` —
+    the trivial trunk body: no hidden layer, no ResBlock stack, no
+    nonlinearity. Ablates whether trunk depth/nonlinearity is earning its
+    parameters, holding everything else (heads, ConditionEncoder,
+    generator, ...) fixed. Composes for free with `router.enabled = true`
+    (gitea #33): a RoutedTrunk of n_experts linear bodies is "mixture of
+    trivial linear experts". `hidden_dim`/`n_blocks`/`dropout`/
+    `block_conditioning` are accepted and ignored, matching
+    `build_expert_body`'s shared factory signature.
+
+    `x` — the trunk's own input (e.g. the noised primary vector for flow
+    matching) — does not already carry conditioning; that's fused in
+    per-body via `cond`. So this concatenates `x` and `cond` itself to
+    remain a valid, conditioning-dependent model.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+        block_conditioning: str = "add",
+    ) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.linear = nn.Linear(in_dim + cond_dim, out_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor | None = None,
+        cond_cat: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.linear(torch.cat([x, cond], dim=-1))
+
+
 def _route_forward(
     experts: nn.ModuleList,
     router: Router,

tests/test_network.py

@@ -10,6 +10,7 @@ from giant.model.network import (
     ConditionEncoder,
     HistoryEncoder,
     MarkovHistory,
+    NoHistory,
     SinusoidalEmbedding,
     Stage1Model,
     Stage2Autoregressive,
@@ -465,16 +466,52 @@ def test_attention_history_step_matches_forward():
     assert torch.allclose(stepped, expected, atol=1e-5)
 
 
+# --- NoHistory (gitea #45) ----------------------------------------------------
+
+
+def test_no_history_shape():
+    hist = NoHistory(in_dim=7, out_dim=12)
+    B, K = 3, 5
+    feat = torch.randn(B, K, 7)
+    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    out = hist(feat, has_prev)
+    assert out.shape == (B, K, 12)
+
+
+def test_no_history_ignores_feat_and_has_prev():
+    hist = NoHistory(in_dim=4, out_dim=6)
+    B, K = 2, 3
+    has_prev_a = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    has_prev_b = torch.zeros(B, K, dtype=torch.bool)
+    feat_a = torch.randn(B, K, 4)
+    feat_b = torch.randn(B, K, 4) * 100
+    out_a = hist(feat_a, has_prev_a)
+    out_b = hist(feat_b, has_prev_b)
+    assert torch.equal(out_a, torch.zeros(B, K, 6))
+    assert torch.equal(out_a, out_b)
+
+
+def test_no_history_uses_base_class_o1_defaults():
+    hist = NoHistory(in_dim=4, out_dim=6)
+    assert hist.init_cache() is None
+    feat = torch.randn(2, 1, 4)
+    has_prev = torch.ones(2, 1, dtype=torch.bool)
+    out, cache = hist.step(feat, has_prev, "unused-cache")
+    assert torch.equal(out, torch.zeros(2, 1, 6))
+    assert cache == "unused-cache"
+
+
 # --- HISTORY_REGISTRY / build_history (gitea #35) ----------------------------
 
 
-def test_history_registry_has_exactly_the_two_known_histories():
-    assert set(HISTORY_REGISTRY) == {"markov", "attention"}
+def test_history_registry_has_exactly_the_known_histories():
+    assert set(HISTORY_REGISTRY) == {"markov", "attention", "none"}
 
 
 def test_build_history_returns_correct_concrete_type():
     assert isinstance(build_history("markov", 4, 6), MarkovHistory)
     assert isinstance(build_history("attention", 4, 8), AttentionHistory)
+    assert isinstance(build_history("none", 4, 6), NoHistory)
 
 
 def test_build_history_unknown_name_raises():
@@ -605,7 +642,7 @@ def test_stage2_autoregressive_n_sec_head_and_type_head_cfg_control_hidden_width
 
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 @pytest.mark.parametrize("generator", ["wgan", "flow"])
-@pytest.mark.parametrize("history", ["markov", "attention"])
+@pytest.mark.parametrize("history", ["markov", "attention", "none"])
 def test_stage2_autoregressive_forward_shape(target, generator, history):
     B, K, emb_dim = 4, 5, 6
     model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K, history=history)

tests/test_router.py

@@ -13,6 +13,8 @@ from giant.model.network import (
     EnergyRouter,
     ExpertTrunk,
     FilmResBlock,
+    LinearTrunk,
+    NoneRouter,
     PdgRouter,
     ProcessRouter,
     ROUTER_REGISTRY,
@@ -73,6 +75,34 @@ def test_energy_router_registered():
     assert ROUTER_REGISTRY["energy"] is EnergyRouter
 
 
+# ── NoneRouter (gitea #45) ──────────────────────────────────────────────────
+
+
+def test_none_router_registered():
+    assert ROUTER_REGISTRY["none"] is NoneRouter
+
+
+def test_none_router_gate_is_uniform():
+    router = NoneRouter(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, torch.full((16, 4), 0.25))
+
+
+def test_none_router_gate_ignores_conditioning():
+    router = NoneRouter(n_experts=3)
+    cond_cont_a, cond_cat_a = _cond(8)
+    cond_cont_b, cond_cat_b = _cond(8)
+    torch.testing.assert_close(router.gate(cond_cont_a, cond_cat_a), router.gate(cond_cont_b, cond_cat_b))
+
+
+def test_none_router_top1_always_expert_zero():
+    router = NoneRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    assert torch.equal(router.top1(cond_cont, cond_cat), torch.zeros(16, dtype=torch.long))
+
+
 def test_energy_router_gate_partition_of_unity():
     router = EnergyRouter(n_experts=4)
     cond_cont, cond_cat = _cond(16)
@@ -168,6 +198,47 @@ def test_build_expert_body_unknown_type_raises():
     raise AssertionError("expected ValueError for unknown trunk type")
 
 
+# ── LinearTrunk (gitea #45) ──────────────────────────────────────────────────
+
+
+def test_trunk_registry_has_linear():
+    assert "linear" in TRUNK_REGISTRY
+    assert TRUNK_REGISTRY["linear"] is LinearTrunk
+
+
+def test_linear_trunk_forward_shape():
+    trunk = build_expert_body("linear", in_dim=9, out_dim=9, hidden_dim=64, n_blocks=6, cond_dim=12)
+    assert trunk.in_dim == 9
+    assert trunk.out_dim == 9
+    x = torch.randn(5, 9)
+    cond = torch.randn(5, 12)
+    out = trunk(x, cond)
+    assert out.shape == (5, 9)
+
+
+def test_linear_trunk_depends_on_x_and_cond():
+    trunk = build_expert_body("linear", in_dim=9, out_dim=9, hidden_dim=64, n_blocks=6, cond_dim=12)
+    x = torch.randn(5, 9)
+    cond_a = torch.randn(5, 12)
+    cond_b = torch.randn(5, 12)
+    assert not torch.allclose(trunk(x, cond_a), trunk(x, cond_b))
+
+
+def test_routed_linear_trunk_is_mixture_of_trivial_experts():
+    """trunk.type = 'linear' composes for free with router.enabled = true
+    (gitea #33's comment on this issue) — a RoutedTrunk of n_experts linear
+    bodies."""
+    router = build_router("energy", n_experts=3)
+    trunk = RoutedTrunk(router, "linear", in_dim=9, out_dim=9, hidden_dim=64, n_res_blocks=6, cond_dim=12)
+    assert len(trunk.experts) == 3
+    assert all(isinstance(e, LinearTrunk) for e in trunk.experts)
+    x = torch.randn(5, 9)
+    cond = torch.randn(5, 12)
+    cond_cont, cond_cat = _cond(5)
+    out = trunk(x, cond, cond_cont, cond_cat)
+    assert out.shape == (5, 9)
+
+
 # ── BLOCK_REGISTRY / build_block (gitea #34) ────────────────────────────────