Implement TrajMixer block

models.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from dataclasses import dataclass
 
 import torch
@@ -14,6 +15,33 @@ from backbones import (
 from targets import PAD_IDX
 
 
+TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v1"
+
+
+def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
+    actual = config.get("model_architecture")
+    if actual != TRAJ_MIXER_ARCHITECTURE:
+        raise ValueError(
+            "This branch only accepts models trained with the TrajMixer "
+            f"architecture marker {TRAJ_MIXER_ARCHITECTURE!r}; got {actual!r}."
+        )
+
+
+def validate_traj_mixer_state_dict(state_dict: Mapping[str, object]) -> None:
+    required_keys = {
+        "blocks.0.mlp.group_align",
+        "blocks.0.mlp.gate_proj",
+        "blocks.0.mlp.value_proj",
+        "blocks.0.mlp.output_proj",
+    }
+    missing = sorted(required_keys.difference(state_dict))
+    if missing:
+        raise ValueError(
+            "Checkpoint is not a TrajMixer checkpoint; missing required "
+            f"parameters: {', '.join(missing)}"
+        )
+
+
 @dataclass
 class DeepHealthOutput:
     hidden: torch.Tensor
@@ -160,6 +188,7 @@ class DeepHealth(nn.Module):
         dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
         extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
+        hidden_group: int = 20,
     ):
         super().__init__()
         if target_mode not in ["next_token", "all_future"]:
@@ -214,6 +243,7 @@ class DeepHealth(nn.Module):
                     use_time_rope=False,
                     use_rbf_bias=False,
                     mlp_dropout=dropout,
+                    hidden_group=hidden_group,
                 ) for _ in range(n_hist_layer)
             ])
             self.rope = None
@@ -227,6 +257,7 @@ class DeepHealth(nn.Module):
                     use_time_rope=True,
                     use_rbf_bias=True,
                     mlp_dropout=dropout,
+                    hidden_group=hidden_group,
                 ) for _ in range(n_hist_layer)
             ])
             self.rope = TimeRoPE(n_embd // n_head)