1"""Fusion of the three resolutions: (a) shared backbone over the concatenated token sequence with resolution2embeddings; (b) separate small per-scale encoders + cross-attention fusion into the 1-minute stream."""3from __future__ import annotations45import torch6import torch.nn as nn78from .encoder import Encoder91011class SharedFusion(nn.Module):12 def __init__(self, d_model: int, n_layers: int, n_heads: int, d_ff: int, dropout: float):13 super().__init__()14 self.enc = Encoder(d_model, n_layers, n_heads, d_ff, dropout)1516 def forward(self, streams: list[tuple[torch.Tensor, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]:17 tok = torch.cat([t for t, _ in streams], dim=1)18 valid = torch.cat([v for _, v in streams], dim=1)19 return self.enc(tok, key_padding_mask=~valid), valid202122class CrossAttentionFusion(nn.Module):23 def __init__(self, d_model: int, n_layers: int, n_heads: int, d_ff: int, dropout: float, n_streams: int = 3):24 super().__init__()25 per = max(1, n_layers // 2)26 self.encs = nn.ModuleList([Encoder(d_model, per, n_heads, d_ff, dropout) for _ in range(n_streams)])27 self.cross = nn.ModuleList([nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)28 for _ in range(n_streams - 1)])29 self.norms = nn.ModuleList([nn.LayerNorm(d_model) for _ in range(n_streams - 1)])30 self.post = Encoder(d_model, max(1, n_layers - per), n_heads, d_ff, dropout)3132 def forward(self, streams: list[tuple[torch.Tensor, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]:33 encoded = [enc(t, key_padding_mask=~v) for enc, (t, v) in zip(self.encs, streams)]34 h, v0 = encoded[0], streams[0][1]35 for k in range(1, len(encoded)):36 kv, vk = encoded[k], streams[k][1]37 kpm = ~vk38 if kpm.all(dim=1).any():39 kpm = kpm.clone(); kpm[kpm.all(dim=1), 0] = False40 a, _ = self.cross[k - 1](self.norms[k - 1](h), kv, kv, key_padding_mask=kpm, need_weights=False)41 h = h + a42 return self.post(h, key_padding_mask=~v0), v043