""" Modular, block-based components for building autoencoders in PyTorch. Core goals: - Composable building blocks with consistent interfaces - Support 2D (B, F) and 3D (B, T, F) tensors where applicable - Simple configs to construct blocks and sequences - Safe-by-default validation and helpful errors This module is intentionally self-contained to allow gradual integration with existing models. It does not mutate current behavior. """ from __future__ import annotations from dataclasses import dataclass from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F # Import config dataclasses that define block configurations try: from .configuration_autoencoder import ( BlockConfig, LinearBlockConfig, AttentionBlockConfig, RecurrentBlockConfig, ConvolutionalBlockConfig, VariationalBlockConfig, ) except Exception: from configuration_autoencoder import ( BlockConfig, LinearBlockConfig, AttentionBlockConfig, RecurrentBlockConfig, ConvolutionalBlockConfig, VariationalBlockConfig, ) # Import shared utilities try: from .utils import _get_activation, _get_norm, _flatten_3d_to_2d, _maybe_restore_3d except Exception: from utils import _get_activation, _get_norm, _flatten_3d_to_2d, _maybe_restore_3d # ---------------------------- Base Block ---------------------------- # class BaseBlock(nn.Module): """Abstract base for all blocks. All blocks should accept 2D (B, F) or 3D (B, T, F) tensors and return the same rank, with last-dim equal to `output_dim`. """ def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: # pragma: no cover - abstract raise NotImplementedError @property def output_dim(self) -> int: # pragma: no cover - abstract raise NotImplementedError # ---------------------------- Residual Base ---------------------------- # class ResidualBlock(BaseBlock): """Base class for blocks supporting residual connections. Implements a safe residual add when input and output dims match; otherwise falls back to a learned projection. Residuals can be scaled. """ def __init__(self, residual: bool = False, residual_scale: float = 1.0, proj_dim_in: Optional[int] = None, proj_dim_out: Optional[int] = None): super().__init__() self.use_residual = residual self.residual_scale = residual_scale self._proj: Optional[nn.Module] = None if residual and proj_dim_in is not None and proj_dim_out is not None and proj_dim_in != proj_dim_out: self._proj = nn.Linear(proj_dim_in, proj_dim_out) def _apply_residual(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: if not self.use_residual: return y x2d, hint = _flatten_3d_to_2d(x) y2d, _ = _flatten_3d_to_2d(y) if x2d.shape[-1] != y2d.shape[-1]: if self._proj is None: self._proj = nn.Linear(x2d.shape[-1], y2d.shape[-1]).to(y2d.device) x2d = self._proj(x2d) out = x2d + self.residual_scale * y2d return _maybe_restore_3d(out, hint) # ---------------------------- LinearBlock ---------------------------- # class LinearBlock(ResidualBlock): """Basic linear transformation with normalization and activation. - Handles both 2D (B, F) and 3D (B, T, F) tensors - Optional normalization: batch|layer|group|instance|none - Configurable activation - Optional dropout - Optional residual connection (with auto projection) """ def __init__(self, cfg: LinearBlockConfig): super().__init__(residual=cfg.use_residual, residual_scale=cfg.residual_scale, proj_dim_in=cfg.input_dim, proj_dim_out=cfg.output_dim) self.cfg = cfg self.linear = nn.Linear(cfg.input_dim, cfg.output_dim) # Normalizations that expect N, C require 2D tensors; for 3D we flatten # For LayerNorm, it supports last-dim directly if cfg.normalization == "layer": self.norm = nn.LayerNorm(cfg.output_dim) else: self.norm = _get_norm(cfg.normalization, cfg.output_dim) self.act = _get_activation(cfg.activation) self.drop = nn.Dropout(cfg.dropout_rate) if cfg.dropout_rate and cfg.dropout_rate > 0 else nn.Identity() @property def output_dim(self) -> int: return self.cfg.output_dim def forward(self, x: torch.Tensor) -> torch.Tensor: x_in = x x2d, hint = _flatten_3d_to_2d(x) y = self.linear(x2d) # Apply norm safely if isinstance(self.norm, (nn.BatchNorm1d, nn.InstanceNorm1d, nn.GroupNorm)): y = self.norm(y) else: # LayerNorm or Identity operates on last dim and supports both 2D/3D; we already have 2D y = self.norm(y) y = self.act(y) y = self.drop(y) y = _maybe_restore_3d(y, hint) return self._apply_residual(x_in, y) # ---------------------------- AttentionBlock ---------------------------- # class AttentionBlock(BaseBlock): """Multi-head self-attention with optional FFN. Expects inputs as 3D (B, T, D) or 2D (B, D) which will be treated as (B, 1, D). Supports optional attn mask and key padding mask via kwargs. """ def __init__(self, cfg: AttentionBlockConfig): super().__init__() self.cfg = cfg d_model = cfg.input_dim self.mha = nn.MultiheadAttention(d_model, num_heads=cfg.num_heads, dropout=cfg.dropout_rate, batch_first=True) self.ln1 = nn.LayerNorm(d_model) ffn_dim = cfg.ffn_dim or (4 * d_model) self.ffn = nn.Sequential( nn.Linear(d_model, ffn_dim), _get_activation("gelu"), nn.Dropout(cfg.dropout_rate), nn.Linear(ffn_dim, d_model), ) self.ln2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(cfg.dropout_rate) @property def output_dim(self) -> int: return self.cfg.input_dim def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None, key_padding_mask: Optional[torch.Tensor] = None) -> torch.Tensor: if x.dim() == 2: x = x.unsqueeze(1) squeeze_back = True else: squeeze_back = False # Self-attention residual = x attn_out, _ = self.mha(x, x, x, attn_mask=attn_mask, key_padding_mask=key_padding_mask, need_weights=False) x = self.ln1(residual + self.dropout(attn_out)) # FFN residual = x x = self.ffn(x) x = self.ln2(residual + self.dropout(x)) if squeeze_back: x = x.squeeze(1) return x # ---------------------------- RecurrentBlock ---------------------------- # class RecurrentBlock(BaseBlock): """RNN processing block supporting LSTM/GRU/RNN. Input: 3D (B, T, F) preferred. If 2D, treated as (B, 1, F). Output dim equals cfg.output_dim if set; otherwise hidden_size * directions. """ def __init__(self, cfg: RecurrentBlockConfig): super().__init__() self.cfg = cfg rnn_type = cfg.rnn_type.lower() rnn_cls = {"lstm": nn.LSTM, "gru": nn.GRU, "rnn": nn.RNN}.get(rnn_type) if rnn_cls is None: raise ValueError(f"Unknown rnn_type: {cfg.rnn_type}") self.rnn = rnn_cls( input_size=cfg.input_dim, hidden_size=cfg.hidden_size, num_layers=cfg.num_layers, batch_first=True, dropout=cfg.dropout_rate if cfg.num_layers > 1 else 0.0, bidirectional=cfg.bidirectional, ) out_dim = cfg.hidden_size * (2 if cfg.bidirectional else 1) self._out_dim = cfg.output_dim or out_dim self.proj = None if self._out_dim == out_dim else nn.Linear(out_dim, self._out_dim) @property def output_dim(self) -> int: return self._out_dim def forward(self, x: torch.Tensor, lengths: Optional[torch.Tensor] = None) -> torch.Tensor: squeeze_back = False if x.dim() == 2: x = x.unsqueeze(1) squeeze_back = True if lengths is not None: x = nn.utils.rnn.pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False) if isinstance(self.rnn, nn.LSTM): out, (h, c) = self.rnn(x) else: out, h = self.rnn(x) if lengths is not None: out, _ = nn.utils.rnn.pad_packed_sequence(out, batch_first=True) # Use last timestep y = out[:, -1, :] if self.proj is not None: y = self.proj(y) if squeeze_back: # Keep 2D output return y # Return (B, 1, D) to keep 3D shape consistent with sequences return y.unsqueeze(1) # ---------------------------- ConvolutionalBlock ---------------------------- # class ConvolutionalBlock(BaseBlock): """1D convolutional block for sequence-like data. Accepts 3D (B, T, F) or 2D (B, F) which is treated as (B, 1, F). """ def __init__(self, cfg: ConvolutionalBlockConfig): super().__init__() self.cfg = cfg # Conv1d expects (B, C_in, L). We interpret features as channels and time as length. # For inputs shaped (B, T, F): we transpose to (B, F, T), apply conv, transpose back. padding = cfg.padding if isinstance(padding, str) and padding == "same": pad = cfg.kernel_size // 2 else: pad = int(padding) self.conv = nn.Conv1d(cfg.input_dim, cfg.output_dim, kernel_size=cfg.kernel_size, padding=pad) # Norm: for Conv1d, use 1d norms over channels if cfg.normalization == "layer": self.norm = nn.GroupNorm(1, cfg.output_dim) # Layer-like over channels else: self.norm = _get_norm(cfg.normalization, cfg.output_dim) self.act = _get_activation(cfg.activation) self.drop = nn.Dropout(cfg.dropout_rate) if cfg.dropout_rate and cfg.dropout_rate > 0 else nn.Identity() @property def output_dim(self) -> int: return self.cfg.output_dim def forward(self, x: torch.Tensor) -> torch.Tensor: squeeze_back = False if x.dim() == 2: x = x.unsqueeze(1) squeeze_back = True # x: (B, T, F) -> (B, F, T) x = x.transpose(1, 2) y = self.conv(x) if isinstance(self.norm, (nn.BatchNorm1d, nn.InstanceNorm1d, nn.GroupNorm)): y = self.norm(y) y = self.act(y) y = self.drop(y) y = y.transpose(1, 2) if squeeze_back: y = y.squeeze(1) return y # ---------------------------- VariationalBlock ---------------------------- # class VariationalBlock(BaseBlock): """Encapsulates mu/logvar projection and reparameterization. Input can be 2D (B, F) or 3D (B, T, F); for 3D, operates per timestep and returns same rank. Stores mu/logvar on the module for downstream loss usage. """ def __init__(self, cfg: VariationalBlockConfig): super().__init__() self.cfg = cfg self.fc_mu = nn.Linear(cfg.input_dim, cfg.latent_dim) self.fc_logvar = nn.Linear(cfg.input_dim, cfg.latent_dim) self._mu: Optional[torch.Tensor] = None self._logvar: Optional[torch.Tensor] = None @property def output_dim(self) -> int: return self.cfg.latent_dim def forward(self, x: torch.Tensor, training: Optional[bool] = None) -> torch.Tensor: if training is None: training = self.training x2d, hint = _flatten_3d_to_2d(x) mu = self.fc_mu(x2d) logvar = self.fc_logvar(x2d) if training: std = torch.exp(0.5 * logvar) eps = torch.randn_like(std) z = mu + eps * std else: z = mu self._mu = mu self._logvar = logvar z = _maybe_restore_3d(z, hint) return z # ---------------------------- BlockSequence ---------------------------- # class BlockSequence(nn.Module): """Compose multiple blocks into a validated sequence. - Validates dimension flow between blocks - Supports gradient checkpointing (per-block) via forward(checkpoint=True) - Supports optional skip connections: pass `skips` as list of (src_idx, dst_idx) """ def __init__(self, blocks: Sequence[BaseBlock], validate_dims: bool = True, skips: Optional[List[Tuple[int, int]]] = None): super().__init__() self.blocks = nn.ModuleList(blocks) self.skips = skips or [] if validate_dims and len(blocks) > 1: for i in range(1, len(blocks)): prev = blocks[i - 1] cur = blocks[i] if getattr(prev, "output_dim", None) is None or getattr(cur, "output_dim", None) is None: continue if prev.output_dim != cur.output_dim and not isinstance(cur, LinearBlock): # Allow LinearBlock to change dims; others must preserve unless they project internally pass # Only warn; users may know what they're doing def forward(self, x: torch.Tensor, checkpoint: bool = False, **kwargs) -> torch.Tensor: activations: Dict[int, torch.Tensor] = {} for i, block in enumerate(self.blocks): if checkpoint and x.requires_grad: x = torch.utils.checkpoint.checkpoint(lambda inp: block(inp, **kwargs), x) else: x = block(x, **kwargs) activations[i] = x # Apply any pending skips to this idx for src, dst in self.skips: if dst == i and src in activations: x = x + activations[src] return x # ---------------------------- Factory ---------------------------- # class BlockFactory: """Factory to build blocks/sequences from configs. This is intentionally minimal; extend as needed. """ @staticmethod def build_block(cfg: Union[BlockConfig, Dict[str, Any]]) -> BaseBlock: # Allow dict-like if isinstance(cfg, dict): type_name = cfg.get("type") # copy and remove 'type' to satisfy dataclass init params = dict(cfg) params.pop("type", None) if type_name == "linear": return LinearBlock(LinearBlockConfig(**params)) if type_name == "attention": return AttentionBlock(AttentionBlockConfig(**params)) if type_name == "recurrent": return RecurrentBlock(RecurrentBlockConfig(**params)) if type_name == "conv1d": return ConvolutionalBlock(ConvolutionalBlockConfig(**params)) raise ValueError(f"Unsupported block type in dict cfg: {type_name} cfg={cfg}") # Dataclass path if isinstance(cfg, LinearBlockConfig) or getattr(cfg, "type", None) == "linear": if not isinstance(cfg, LinearBlockConfig): cfg = LinearBlockConfig(**cfg.__dict__) # type: ignore[arg-type] return LinearBlock(cfg) if isinstance(cfg, AttentionBlockConfig) or getattr(cfg, "type", None) == "attention": if not isinstance(cfg, AttentionBlockConfig): cfg = AttentionBlockConfig(**cfg.__dict__) # type: ignore[arg-type] return AttentionBlock(cfg) if isinstance(cfg, RecurrentBlockConfig) or getattr(cfg, "type", None) == "recurrent": if not isinstance(cfg, RecurrentBlockConfig): cfg = RecurrentBlockConfig(**cfg.__dict__) # type: ignore[arg-type] return RecurrentBlock(cfg) if isinstance(cfg, ConvolutionalBlockConfig) or getattr(cfg, "type", None) == "conv1d": if not isinstance(cfg, ConvolutionalBlockConfig): cfg = ConvolutionalBlockConfig(**cfg.__dict__) # type: ignore[arg-type] return ConvolutionalBlock(cfg) if isinstance(cfg, VariationalBlockConfig) or getattr(cfg, "type", None) == "variational": if not isinstance(cfg, VariationalBlockConfig): cfg = VariationalBlockConfig(**cfg.__dict__) # type: ignore[arg-type] return VariationalBlock(cfg) raise ValueError(f"Unsupported block type: {cfg}") @staticmethod def build_sequence(configs: Sequence[Union[BlockConfig, Dict[str, Any]]]) -> BlockSequence: blocks: List[BaseBlock] = [BlockFactory.build_block(c) for c in configs] return BlockSequence(blocks) __all__ = [ "BlockConfig", "LinearBlockConfig", "AttentionBlockConfig", "RecurrentBlockConfig", "ConvolutionalBlockConfig", "VariationalBlockConfig", "BaseBlock", "ResidualBlock", "LinearBlock", "AttentionBlock", "RecurrentBlock", "ConvolutionalBlock", "VariationalBlock", "BlockSequence", "BlockFactory", ]