1"""History-only regime conditioning: trailing volatility / relative volume statistics -> embedding added to the2decoder queries (option 'stats'); 'learned' = a small codebook selected by a soft classifier over the same stats;3'none' = disabled (ablation H3)."""4from __future__ import annotations56import torch7import torch.nn as nn8910class RegimeEmbedding(nn.Module):11 def __init__(self, f_global: int, d_model: int, mode: str = "stats", n_codes: int = 8):12 super().__init__()13 self.mode = mode14 if mode == "stats":15 self.mlp = nn.Sequential(nn.Linear(f_global, d_model), nn.GELU(), nn.Linear(d_model, d_model))16 elif mode == "learned":17 self.cls = nn.Sequential(nn.Linear(f_global, d_model), nn.GELU(), nn.Linear(d_model, n_codes))18 self.codes = nn.Embedding(n_codes, d_model)19 elif mode != "none":20 raise ValueError(mode)2122 def forward(self, glob: torch.Tensor) -> torch.Tensor | None:23 if self.mode == "none":24 return None25 if self.mode == "stats":26 return self.mlp(glob)27 w = torch.softmax(self.cls(glob), dim=-1)28 return w @ self.codes.weight29