import torch import torch.nn as nn import torch.nn.functional as F import math from typing import Dict, Any, Optional class PatchEmbedding(nn.Module): def __init__(self, d_model: int, frames_per_patch: int = 6, expansion: int = 2): super().__init__() self.d_model = d_model self.frames_per_patch = frames_per_patch # Frame embedding (collapse pitch dim) self.conv2d = nn.Conv2d( in_channels=1, out_channels=d_model, kernel_size=(88, 1), stride=(1, 1), padding=(0, 0), ) self.norm_frame = nn.LayerNorm(d_model) # anti-aliasing conv on time axis self.aa = nn.Conv1d(d_model, d_model, kernel_size=3, stride=1, padding=1, groups=d_model, bias=False) # Late temporal pooling (downsample frames -> patches) self.glu_conv = nn.Conv1d( in_channels=d_model, out_channels=d_model * expansion * 2, kernel_size=frames_per_patch, stride=frames_per_patch, padding=0, bias=True, ) self.project = nn.Conv1d( in_channels=d_model * expansion, out_channels=d_model, kernel_size=1, ) self.norm_temporal = nn.LayerNorm(d_model) def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, 1, 88, T) x = self.conv2d(x) # (B, C, 1, T) x = x.squeeze(2).transpose(1, 2) # (B, T, C) x = self.norm_frame(x) # anti-aliased and temporal pooling x = x.transpose(1, 2) # (B, C, T) x = self.aa(x) # (B, C, T) v, g = self.glu_conv(x).chunk(2, dim=1) x = self.project(v * torch.sigmoid(g)) # (B, C, T//k) x = x.transpose(1, 2) # (B, T//k, C) return self.norm_temporal(x) class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout=0.1, max_len=5000): super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(p=dropout) pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0).transpose(0, 1) self.register_buffer('pe', pe) def forward(self, x): x = x.transpose(0, 1) x = x + self.pe[:x.size(0), :] x = self.dropout(x) return x.transpose(0, 1) class RelativeTransformerEncoderLayer(nn.Module): def __init__(self, d_model: int, nhead: int, dim_feedforward: int, dropout: float, activation: str = 'gelu'): super().__init__() if d_model % nhead != 0: raise ValueError("d_model must be divisible by nhead.") self.d_model = d_model self.nhead = nhead self.head_dim = d_model // nhead self.qkv_proj = nn.Linear(d_model, 3 * d_model) self.attn_dropout = nn.Dropout(dropout) self.out_proj = nn.Linear(d_model, d_model) self.resid_dropout = nn.Dropout(dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.linear1 = nn.Linear(d_model, dim_feedforward) self.linear2 = nn.Linear(dim_feedforward, d_model) self.ff_dropout = nn.Dropout(dropout) if activation == 'gelu': self.activation_fn = F.gelu elif activation == 'relu': self.activation_fn = F.relu else: raise ValueError(f"Unsupported activation: {activation}") def forward( self, src: torch.Tensor, src_mask: Optional[torch.Tensor] = None, src_key_padding_mask: Optional[torch.Tensor] = None, attn_bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: # src: (B, T_new, C) bsz, seq_len_new, _ = src.size() qkv = self.qkv_proj(src) # (B, T_new, 3*C) q, k, v = qkv.chunk(3, dim=-1) q = q.view(bsz, seq_len_new, self.nhead, self.head_dim) k_all = k.view(bsz, seq_len_new, self.nhead, self.head_dim) v_all = v.view(bsz, seq_len_new, self.nhead, self.head_dim) # Attention: queries are only for current tokens; keys include past+current attn_scores = torch.einsum('bthd,bshd->bhts', q, k_all) / math.sqrt(self.head_dim) # (B, H, T_new, T_total) # Additive or boolean mask over attention logits if src_mask is not None: if src_mask.dtype == torch.bool: attn_scores = attn_scores.masked_fill(src_mask.unsqueeze(0), float('-inf')) else: attn_scores = attn_scores + src_mask.unsqueeze(0) # Key padding mask if src_key_padding_mask is not None: key_mask = src_key_padding_mask.unsqueeze(1).unsqueeze(2) # (B,1,1,T) attn_scores = attn_scores.masked_fill(key_mask, float('-inf')) if attn_bias is not None: # Support 3D (H, T, T) or 4D (B, H, T, T) if attn_bias.dim() == 3: attn_scores = attn_scores + attn_bias.unsqueeze(0) elif attn_bias.dim() == 4: attn_scores = attn_scores + attn_bias else: raise ValueError("attn_bias must be 3D or 4D tensor if provided") attn_weights = F.softmax(attn_scores, dim=-1) attn_weights = self.attn_dropout(attn_weights) context = torch.einsum('bhts,bshd->bthd', attn_weights, v_all) # (B, T_new, H, D) context = context.contiguous().view(bsz, seq_len_new, self.d_model) # (B, T_new, C) attn_out = self.out_proj(context) src = src + self.resid_dropout(attn_out) src = self.norm1(src) ff = self.linear2(self.ff_dropout(self.activation_fn(self.linear1(src)))) src = src + self.resid_dropout(ff) src = self.norm2(src) return src def downsample_key_padding_mask(mask: torch.Tensor, frames_per_patch: int) -> torch.Tensor: # mask: (B, T) where True denotes padding. bsz, total_len = mask.shape if total_len < frames_per_patch: # No valid output tokens from temporal pooling return mask.new_ones((bsz, 0), dtype=mask.dtype) out_len = total_len // frames_per_patch trimmed = mask[:, :out_len * frames_per_patch] grouped = trimmed.view(bsz, out_len, frames_per_patch) return grouped.all(dim=-1) class RelativePositionBias(nn.Module): def __init__(self, num_heads: int, max_distance: int): super().__init__() if max_distance < 1: raise ValueError("max_distance must be >= 1") self.max_distance = max_distance self.num_heads = num_heads # Table over relative distances in [-max_distance+1, max_distance-1] self.bias = nn.Parameter(torch.zeros(2 * max_distance - 1, num_heads)) def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: # Compute clipped relative position indices pos = torch.arange(seq_len, device=device) rel = pos[:, None] - pos[None, :] # (T, T) rel = rel.clamp(-self.max_distance + 1, self.max_distance - 1) rel = rel + self.max_distance - 1 # shift to [0, 2*max_distance-2] bias = self.bias[rel] # (T, T, H) return bias.permute(2, 0, 1).to(dtype=dtype) # (H, T, T) class RelativeTransformerEncoder(nn.Module): def __init__(self, num_layers: int, d_model: int, nhead: int, dim_feedforward: int, dropout: float, activation: str = 'gelu', relative_position_bias = None): super().__init__() self.layers = nn.ModuleList([ RelativeTransformerEncoderLayer( d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, activation=activation, ) for _ in range(num_layers) ]) self.norm = nn.LayerNorm(d_model) self.rpb = relative_position_bias self.nhead = nhead def forward( self, src: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: output = src if self.rpb is not None: attn_bias = self.rpb(src.size(1), device=src.device, dtype=src.dtype) else: attn_bias = None for mod in self.layers: output = mod( output, src_key_padding_mask=src_key_padding_mask, attn_bias=attn_bias, ) output = self.norm(output) return output class ChordDecomposeProjection(nn.Module): def __init__(self, d_model: int, vocab_sizes: Dict[str, int]): super().__init__() self.d_model = d_model self.vocab_sizes = vocab_sizes self.boundary_head = nn.Sequential( nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Linear(d_model // 2, 1), ) self.projection_heads = nn.ModuleDict() for comp, size in self.vocab_sizes.items(): self.projection_heads[comp] = nn.Sequential( nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Linear(d_model // 2, size), ) def forward(self, x: torch.Tensor) -> torch.Tensor: boundary_logits = self.boundary_head(x) output = {'boundary': boundary_logits.squeeze(-1)} for comp, head in self.projection_heads.items(): output[comp] = head(x) return output class ChordRecognitionModel(nn.Module): def __init__(self, model_config: Dict[str, Any], vocab_sizes: Dict[str, int]): super().__init__() self.config = model_config self.vocab_sizes = vocab_sizes self.d_model = self.config['d_model'] # Encoder: shared patch embedding and relative transformer (unchanged) self.embedding = PatchEmbedding( d_model=self.d_model, frames_per_patch=self.config['frames_per_patch'], expansion=2, ) self.input_dropout = nn.Dropout(self.config['dropout']) rpb = RelativePositionBias( num_heads=self.config['n_head'], max_distance=self.config['n_beats'] * self.config['label_resolution'] ) self.relative_transformer_encoder = RelativeTransformerEncoder( num_layers=self.config['num_encoder_layers'], d_model=self.d_model, nhead=self.config['n_head'], dim_feedforward=self.config['dim_feedforward'], dropout=self.config['dropout'], activation='gelu', relative_position_bias=rpb, ) # Boundary head, smoother, and FiLM gating d_b = max(1, self.d_model // 4) k_b = int(self.config.get('boundary_kernel', 5)) self.boundary_head = nn.Linear(self.d_model, 1) self.boundary_smoother = nn.Conv1d( in_channels=1, out_channels=1, kernel_size=k_b, padding=k_b // 2, groups=1, bias=True, ) self.boundary_e0 = nn.Parameter(torch.zeros(d_b)) self.boundary_e1 = nn.Parameter(torch.randn(d_b) * 0.02) # Optional key context (3-part setting). Infer size from vocab_sizes if provided self.Vq = int(vocab_sizes.get('quality', 0)) self.Vr = int(vocab_sizes.get('root', 0)) self.Vb = int(vocab_sizes.get('bass', 0)) # FiLM layers take boundary embedding self.film_ln_in = nn.LayerNorm(self.d_model + d_b) self.film_ln_h = nn.LayerNorm(self.d_model) self.film_mlp = nn.Linear(self.d_model + d_b, 2 * self.d_model) # Triple-token decoder: embeddings and heads self.mask_id_q = int(self.config.get('mask_id_q', self.Vq)) self.mask_id_r = int(self.config.get('mask_id_r', self.Vr)) self.mask_id_b = int(self.config.get('mask_id_b', self.Vb)) self.emb_q = nn.Embedding(self.Vq + 1, self.d_model) self.emb_r = nn.Embedding(self.Vr + 1, self.d_model) self.emb_b = nn.Embedding(self.Vb + 1, self.d_model) dec_heads = int(self.config.get('dec_heads', 4)) dec_mlp_ratio = int(self.config.get('dec_mlp_ratio', 4)) dec_layers = int(self.config.get('dec_layers', 1)) dec_dropout = float(self.config.get('dec_dropout', 0.1)) self.window_radius = int(self.config.get('window_radius', 2)) self.decoder_layers = nn.ModuleList([ KTokenDecoderLayer( d_model=self.d_model, nhead=dec_heads, mlp_ratio=dec_mlp_ratio, dropout=dec_dropout, ) for _ in range(dec_layers) ]) self.dec_norm = nn.LayerNorm(self.d_model) self.head_q = nn.Linear(self.d_model, self.Vq) self.head_r = nn.Linear(self.d_model, self.Vr) self.head_b = nn.Linear(self.d_model, self.Vb) # Legacy decompose projection for compatibility when training in decompose mode self.chord_decompose_projection = ChordDecomposeProjection(self.d_model, self.vocab_sizes) def forward(self, encoder_input: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: H, _ = self._encode(encoder_input, src_key_padding_mask) out = self.chord_decompose_projection(H) return out # ===== Utilities ===== def _encode(self, encoder_input: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor]) -> (torch.Tensor, torch.Tensor): x = self.embedding(encoder_input) mask_down = None if src_key_padding_mask is not None: mask_down = downsample_key_padding_mask(src_key_padding_mask, self.config['frames_per_patch']) x = self.input_dropout(x) H = self.relative_transformer_encoder(x, src_key_padding_mask=mask_down) boundary_logits = self.boundary_head(H).squeeze(-1) return H, boundary_logits def _smooth_boundary(self, boundary_logits: torch.Tensor) -> torch.Tensor: b = boundary_logits.unsqueeze(1) smoothed = self.boundary_smoother(b) return torch.sigmoid(smoothed.squeeze(1)) def _apply_film(self, H: torch.Tensor, b_soft: torch.Tensor) -> torch.Tensor: B, T, D = H.shape e0 = self.boundary_e0.view(1, 1, -1).expand(B, T, -1) e1 = self.boundary_e1.view(1, 1, -1).expand(B, T, -1) # soft embedding for boundary eb = b_soft.unsqueeze(-1) * e1 + (1.0 - b_soft).unsqueeze(-1) * e0 film_in = torch.cat([H, eb], dim=-1) # layer norm and linear projection film_in = self.film_ln_in(film_in) gamma, beta = self.film_mlp(film_in).chunk(2, dim=-1) Z = self.film_ln_h(H) * (1.0 + gamma) + beta return Z def _build_local_windows(self, H: torch.Tensor, radius: int) -> torch.Tensor: x = H.transpose(1, 2) padded = F.pad(x, (radius, radius), mode='replicate') win = padded.unfold(dimension=2, size=2 * radius + 1, step=1) win = win.permute(0, 2, 3, 1).contiguous() return win def _build_context(self, H: torch.Tensor, Z: torch.Tensor, b_soft: torch.Tensor) -> torch.Tensor: local = self._build_local_windows(H, self.window_radius) z = Z.unsqueeze(2) parts = [z, local] C = torch.cat(parts, dim=2) return C def _embed_tokens(self, ids_q: torch.Tensor, ids_r: torch.Tensor, ids_b: torch.Tensor) -> torch.Tensor: xq = self.emb_q(ids_q) xr = self.emb_r(ids_r) xb = self.emb_b(ids_b) X = torch.stack([xq, xr, xb], dim=2) return X def _run_decoder(self, X: torch.Tensor, C: torch.Tensor): x = X for layer in self.decoder_layers: x = layer(x, C) x = self.dec_norm(x) xq = x[:, :, 0, :] xr = x[:, :, 1, :] xb = x[:, :, 2, :] logits_q = self.head_q(xq) logits_r = self.head_r(xr) logits_b = self.head_b(xb) return logits_q, logits_r, logits_b # ===== Training forward ===== def forward_train(self, encoder_input: torch.Tensor, targets: Dict[str, torch.Tensor], src_key_padding_mask: Optional[torch.Tensor] = None, target_mask: Optional[torch.Tensor] = None) -> Dict[str, Any]: device = encoder_input.device H, boundary_logits = self._encode(encoder_input, src_key_padding_mask) # No key prediction/context in training b_soft = self._smooth_boundary(boundary_logits) Z = self._apply_film(H, b_soft) C = self._build_context(H, Z, b_soft) tgt_q = targets['quality'] tgt_r = targets['root'] tgt_b = targets['bass'] B, T = tgt_q.shape if target_mask is None: target_mask = torch.ones(B, T, dtype=torch.bool, device=device) # mask n slots randomly per (B,T) across 3 slots [q,r,b] k_rand = torch.randint(1, 4, (B, T), device=device) rand_scores = torch.rand(B, T, 3, device=device) top_vals, top_idx = torch.topk(rand_scores, k=3, dim=-1) mask_slots = torch.zeros(B, T, 3, dtype=torch.bool, device=device) # enable first k indices per position for kk in range(1, 4): sel = (k_rand == kk) if sel.any(): idx_sel = top_idx[sel][:, :kk] row = mask_slots[sel] if idx_sel.numel() > 0: row.scatter_(dim=1, index=idx_sel, value=True) mask_slots[sel] = row ids_q = tgt_q.clone() ids_r = tgt_r.clone() ids_b = tgt_b.clone() ids_q[mask_slots[:, :, 0]] = self.mask_id_q ids_r[mask_slots[:, :, 1]] = self.mask_id_r ids_b[mask_slots[:, :, 2]] = self.mask_id_b X = self._embed_tokens(ids_q, ids_r, ids_b) logits_q, logits_r, logits_b = self._run_decoder(X, C) def ce_masked(logits: Optional[torch.Tensor], target: torch.Tensor, slot_mask: torch.Tensor) -> torch.Tensor: # Build supervision mask and safe targets to avoid CUDA asserts from out-of-range labels m = slot_mask & target_mask # (B,T) supervised locations num_classes = logits.size(-1) safe_target = torch.where( m, target.clamp(min=0, max=num_classes - 1), torch.zeros_like(target) ) ce = F.cross_entropy(logits.transpose(1, 2), safe_target, reduction='none') denom = m.float().sum().clamp(min=1.0) return (ce * m.float()).sum() / denom loss_q = ce_masked(logits_q, tgt_q, mask_slots[:, :, 0]) loss_r = ce_masked(logits_r, tgt_r, mask_slots[:, :, 1]) loss_b = ce_masked(logits_b, tgt_b, mask_slots[:, :, 2]) bce = F.binary_cross_entropy_with_logits(boundary_logits, targets['boundary'].to(boundary_logits.dtype), pos_weight=torch.tensor(2.0, device=device), reduction='none') loss_boundary = (bce * target_mask.float()).sum() / target_mask.float().sum().clamp(min=1.0) total_loss = loss_q + loss_r + loss_b + loss_boundary * 3 with torch.no_grad(): stats = {} for name, logits, target, m in [ ('quality', logits_q, tgt_q, mask_slots[:, :, 0]), ('root', logits_r, tgt_r, mask_slots[:, :, 1]), ('bass', logits_b, tgt_b, mask_slots[:, :, 2]), ]: if logits is None: stats[f'acc_{name}'] = 0.0 stats[f'conf_{name}'] = 0.0 stats[f'ece_{name}'] = 0.0 else: pred = logits.argmax(dim=-1) sel = (m & target_mask) denom = sel.float().sum().clamp(min=1.0) acc = (pred[sel] == target[sel]).float().sum() / denom prob = logits.float().softmax(dim=-1) conf = prob.max(dim=-1).values mean_conf = conf[sel].sum() / denom # simple ECE ece = torch.tensor(0.0, device=device) bins = torch.linspace(0, 1, steps=11, device=device) conf_flat = conf[sel] pred_flat = pred[sel] tgt_flat = target[sel] for i in range(10): lo, hi = bins[i], bins[i+1] mask_bin = (conf_flat >= lo) & (conf_flat < hi if i < 9 else conf_flat <= hi) if mask_bin.sum() > 0: acc_bin = (pred_flat[mask_bin] == tgt_flat[mask_bin]).float().mean() conf_bin = conf_flat[mask_bin].mean() ece = ece + (mask_bin.float().mean() * (acc_bin - conf_bin).abs()) stats[f'acc_{name}'] = acc.item() stats[f'conf_{name}'] = mean_conf.item() stats[f'ece_{name}'] = ece.item() return { 'loss': total_loss, 'loss_map': { 'quality': loss_q, 'root': loss_r, 'bass': loss_b, 'boundary': loss_boundary, }, 'logits': { 'quality': logits_q, 'root': logits_r, 'bass': logits_b, }, 'mask_slots': mask_slots, # (B,T,3) bool in order [q,r,b] 'boundary_logits': boundary_logits, 'stats': stats, } # ===== Inference forward ===== def forward_infer(self, encoder_input: torch.Tensor, src_key_padding_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]: device = encoder_input.device H, boundary_logits = self._encode(encoder_input, src_key_padding_mask) # No key prediction/context in inference b_soft = self._smooth_boundary(boundary_logits) Z = self._apply_film(H, b_soft) C = self._build_context(H, Z, b_soft) B, T, _ = H.shape ids_q = torch.full((B, T), self.mask_id_q, dtype=torch.long, device=device) ids_r = torch.full((B, T), self.mask_id_r, dtype=torch.long, device=device) ids_b = torch.full((B, T), self.mask_id_b, dtype=torch.long, device=device) filled_q = torch.zeros((B, T), dtype=torch.bool, device=device) filled_r = torch.zeros((B, T), dtype=torch.bool, device=device) filled_b = torch.zeros((B, T), dtype=torch.bool, device=device) # Track decode order per time step: 0=quality, 1=root, 2=bass decode_order = torch.full((B, T, 3), -1, dtype=torch.long, device=device) order_pos = 0 for step in (3, 2, 1): X = self._embed_tokens(ids_q, ids_r, ids_b) logits_q, logits_r, logits_b = self._run_decoder(X, C) pq = logits_q.softmax(dim=-1) pr = logits_r.softmax(dim=-1) pb = logits_b.softmax(dim=-1) conf_q = pq.max(dim=-1).values conf_r = pr.max(dim=-1).values conf_b = pb.max(dim=-1).values conf_q = conf_q.masked_fill(filled_q, float('-inf')) conf_r = conf_r.masked_fill(filled_r, float('-inf')) conf_b = conf_b.masked_fill(filled_b, float('-inf')) conf = torch.stack([conf_q, conf_r, conf_b], dim=-1) take_slot = conf.argmax(dim=-1) # record order at this step decode_order[:, :, order_pos] = take_slot order_pos += 1 pred_q = logits_q.argmax(dim=-1) commit_q = (take_slot == 0) | ((step == 1) & (~filled_q)) ids_q[commit_q] = pred_q[commit_q] filled_q = filled_q | commit_q pred_r = logits_r.argmax(dim=-1) commit_r = (take_slot == 1) | ((step == 1) & (~filled_r)) ids_r[commit_r] = pred_r[commit_r] filled_r = filled_r | commit_r pred_b = logits_b.argmax(dim=-1) commit_b = (take_slot == 2) | ((step == 1) & (~filled_b)) ids_b[commit_b] = pred_b[commit_b] filled_b = filled_b | commit_b # final confidences X = self._embed_tokens(ids_q, ids_r, ids_b) logits_q, logits_r, logits_b = self._run_decoder(X, C) conf_q = logits_q.softmax(dim=-1).max(dim=-1).values conf_r = logits_r.softmax(dim=-1).max(dim=-1).values conf_b = logits_b.softmax(dim=-1).max(dim=-1).values return { 'quality': ids_q, 'root': ids_r, 'bass': ids_b, 'conf_quality': conf_q, 'conf_root': conf_r, 'conf_bass': conf_b, 'boundary': boundary_logits, 'decode_order': decode_order, } class KTokenDecoderLayer(nn.Module): def __init__(self, d_model: int, nhead: int, mlp_ratio: int, dropout: float): super().__init__() self.d_model = d_model self.nhead = nhead self.head_dim = d_model // nhead if self.head_dim * nhead != d_model: raise ValueError("d_model must be divisible by nhead") # self-attention over K tokens self.sa_qkv = nn.Linear(d_model, 3 * d_model) self.sa_out = nn.Linear(d_model, d_model) self.sa_ln = nn.LayerNorm(d_model) self.sa_drop = nn.Dropout(dropout) # cross-attention to context self.ca_q = nn.Linear(d_model, d_model) self.ca_kv = nn.Linear(d_model, 2 * d_model) self.ca_out = nn.Linear(d_model, d_model) self.ca_ln = nn.LayerNorm(d_model) self.ca_drop = nn.Dropout(dropout) # ffn hidden = d_model * mlp_ratio self.ff_ln = nn.LayerNorm(d_model) self.ff = nn.Sequential( nn.Linear(d_model, hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden, d_model), ) self.ff_drop = nn.Dropout(dropout) def _attn(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: # q,k,v: (N, L, H, D) attn = torch.einsum('nlhd,nshd->nhls', q, k) / math.sqrt(q.size(-1)) # (N,H,L,S) attn = torch.softmax(attn, dim=-1) ctx = torch.einsum('nhls,nshd->nlhd', attn, v) # (N,L,H,D) ctx = ctx.contiguous().view(q.size(0), q.size(1), -1) # (N,L,C) return ctx def forward(self, X: torch.Tensor, C: torch.Tensor) -> torch.Tensor: # X: (B,T,K,D), C: (B,T,Lc,D) B, T, K, D = X.shape Lc = C.size(2) N = B * T # reshape x = X.view(N, K, D) c = C.view(N, Lc, D) # self-attn (over K) x_norm = self.sa_ln(x) qkv = self.sa_qkv(x_norm) q, k, v = qkv.chunk(3, dim=-1) q = q.view(N, K, self.nhead, self.head_dim) k = k.view(N, K, self.nhead, self.head_dim) v = v.view(N, K, self.nhead, self.head_dim) sa_ctx = self._attn(q, k, v) # (N,K,C) x = x + self.sa_drop(self.sa_out(sa_ctx)) # cross-attn (queries = tokens, keys/values = context) x_norm = self.ca_ln(x) q = self.ca_q(x_norm).view(N, K, self.nhead, self.head_dim) kv = self.ca_kv(c) k, v = kv.chunk(2, dim=-1) k = k.view(N, Lc, self.nhead, self.head_dim) v = v.view(N, Lc, self.nhead, self.head_dim) ca_ctx = self._attn(q, k, v) # (N,K,C) x = x + self.ca_drop(self.ca_out(ca_ctx)) # ffn x_norm = self.ff_ln(x) x = x + self.ff_drop(self.ff(x_norm)) return x.view(B, T, K, D)