model.py

"""
TRM (Tiny Recursive Model) Core Architecture Components

Implementation of the TRM model as described in:
"Less is More: Recursive Reasoning with Tiny Networks" (arXiv:2510.04871v1)
Author: Alexia Jolicoeur-Martineau

This module implements the fundamental building blocks:
- RMSNorm: Root Mean Square Layer Normalization
- RotaryEmbedding: Rotary Position Embeddings (RoPE)
- SwiGLU: SwiGLU activation function
- MLPMixer: MLP applied on sequence dimension
- NetworkBlock: 2-layer network with sequence and channel mixing
- InputEmbedding: Token + position embedding
- OutputHead: Projection to vocabulary
- QHead: Halting probability prediction

All hyperparameters are hardcoded as specified in the paper (Page 11):
- hidden_size = 512
- context_length = 81 (9x9 Sudoku grid)
- vocab_size = 11 (tokens 0-10: PAD + 0-9)
- num_layers = 2
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def _find_multiple(a, b):
    """Find the smallest multiple of b that is >= a.

    Used for rounding intermediate dimensions to multiples of specific values
    (e.g., 256) for better hardware utilization and memory alignment.

    Args:
        a: The minimum value to round up from
        b: The multiple to round to

    Returns:
        Smallest multiple of b that is >= a

    Reference:
        From TinyRecursiveModels reference implementation
    """
    return (-(a // -b)) * b

def trunc_normal_init_(tensor: torch.Tensor, std: float = 1.0, lower: float = -2.0, upper: float = 2.0):
    # NOTE: PyTorch nn.init.trunc_normal_ is not mathematically correct, the std dev is not actually the std dev of initialized tensor
    # This function is a PyTorch version of jax truncated normal init (default init method in flax)
    # https://github.com/jax-ml/jax/blob/main/jax/_src/random.py#L807-L848
    # https://github.com/jax-ml/jax/blob/main/jax/_src/nn/initializers.py#L162-L199

    with torch.no_grad():
        if std == 0:
            tensor.zero_()
        else:
            sqrt2 = math.sqrt(2)
            a = math.erf(lower / sqrt2)
            b = math.erf(upper / sqrt2)
            z = (b - a) / 2

            c = (2 * math.pi) ** -0.5
            pdf_u = c * math.exp(-0.5 * lower ** 2)
            pdf_l = c * math.exp(-0.5 * upper ** 2)
            comp_std = std / math.sqrt(1 - (upper * pdf_u - lower * pdf_l) / z - ((pdf_u - pdf_l) / z) ** 2)

            tensor.uniform_(a, b)
            tensor.erfinv_()
            tensor.mul_(sqrt2 * comp_std)
            tensor.clip_(lower * comp_std, upper * comp_std)

    return tensor


class CastedLinear(nn.Module):
    """
    Self-initializing Linear layer with LeCun normal initialization.

    Drop-in replacement for nn.Linear that automatically applies proper initialization:
    - Weight: truncated normal with std = 1/sqrt(in_features) (LeCun fan-in scaling)
    - Bias: zero initialization

    Also provides automatic dtype casting in forward pass for mixed precision training.

    Reference:
        From TinyRecursiveModels/models/layers.py (lines 44-60)
        Used throughout TRM to eliminate repetitive initialization code
    """

    def __init__(self, in_features: int, out_features: int, bias: bool):
        """Initialize CastedLinear with automatic LeCun normal initialization.

        Args:
            in_features: Input dimension (used for LeCun scaling: std = 1/sqrt(in_features))
            out_features: Output dimension
            bias: Whether to include bias term (auto-initialized to zero if True)
        """
        super().__init__()

        # Weight initialized with truncated LeCun normal (std = 1/sqrt(fan_in))
        self.weight = nn.Parameter(
            trunc_normal_init_(
                torch.empty((out_features, in_features)),
                std=1.0 / (in_features ** 0.5)
            )
        )

        # Bias initialized to zero if enabled
        self.bias = None
        if bias:
            self.bias = nn.Parameter(torch.zeros((out_features,)))

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """Forward pass with automatic dtype casting for mixed precision.

        Args:
            input: Input tensor of shape [..., in_features]

        Returns:
            Output tensor of shape [..., out_features]
        """
        # Cast weight and bias to match input dtype for mixed precision support
        weight = self.weight.to(input.dtype)
        bias = self.bias.to(input.dtype) if self.bias is not None else None
        return F.linear(input, weight, bias=bias)


def rms_norm(hidden_states: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
    """
    Functional RMSNorm matching reference implementation.

    Root Mean Square Layer Normalization without learnable parameters.
    Casts to float32 for numerical stability, then back to input dtype.

    As specified in Section 2.1 of the paper (from HRM) and matching
    the reference TinyRecursiveModels implementation.

    Args:
        hidden_states: Input tensor of shape [..., dim]
        eps: Small constant for numerical stability (default: 1e-5)

    Returns:
        Normalized tensor of same shape and dtype as input

    Reference:
        TinyRecursiveModels/models/layers.py lines 163-169
    """
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)

    variance = hidden_states.square().mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + eps)

    return hidden_states.to(input_dtype)


class RotaryEmbedding(nn.Module):
    """
    Rotary Position Embeddings (RoPE).

    As specified in Section 2.1 of the paper (from HRM).
    Encodes position information through rotation in complex plane.
    """

    def __init__(self, dim: int, base: int = 10000):
        """
        Args:
            dim: Feature dimension (must be even)
            base: Base for frequency computation
        """
        super().__init__()
        self.dim = dim

        # Compute frequency for each dimension pair
        # inv_freq shape: [dim // 2]
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Apply rotary embeddings to input.

        Args:
            x: Input tensor of shape [batch, seq_len, dim]

        Returns:
            Tensor with rotary embeddings applied, same shape as input
        """
        batch_size, seq_len, _ = x.shape

        # Create position indices: [seq_len]
        position = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)

        # Compute frequencies: [seq_len, dim // 2]
        freqs = torch.outer(position, self.inv_freq)

        # Compute cos and sin for each frequency
        # Shape: [seq_len, dim // 2]
        cos_freqs = freqs.cos()
        sin_freqs = freqs.sin()

        # Rotate pairs: for each pair (x[2i], x[2i+1]), apply 2D rotation
        # Split into even and odd indices
        x_even = x[..., 0::2]  # [batch, seq_len, dim//2]
        x_odd = x[..., 1::2]   # [batch, seq_len, dim//2]

        # Apply rotation
        # [cos*x_even - sin*x_odd, sin*x_even + cos*x_odd]
        cos_part = cos_freqs.unsqueeze(0)  # [1, seq_len, dim//2]
        sin_part = sin_freqs.unsqueeze(0)  # [1, seq_len, dim//2]

        rotated_even = x_even * cos_part - x_odd * sin_part
        rotated_odd = x_even * sin_part + x_odd * cos_part

        # Interleave back
        rotated = torch.zeros_like(x)
        rotated[..., 0::2] = rotated_even
        rotated[..., 1::2] = rotated_odd

        return rotated


class SwiGLU(nn.Module):
    """
    SwiGLU activation function with fused gate/up projection.

    As specified in Section 2.1 of the paper (from HRM).
    Combines Swish (SiLU) activation with Gated Linear Unit.

    This implementation uses a fused projection pattern for memory efficiency:
    - Single projection to 2*intermediate_dim, then chunk into gate and up
    - More efficient than separate projections (reduces memory allocations)

    Formula: SwiGLU(x) = down_proj(silu(gate) ⊙ up)
    where gate, up = gate_up_proj(x).chunk(2)

    Reference:
        From TinyRecursiveModels reference implementation
        Common optimization in LLaMA and other modern transformers
    """

    def __init__(self, hidden_size: int, expansion: float):
        """
        Args:
            hidden_size: Input/output dimension
            expansion: Expansion factor for intermediate dimension
                      Calculates intermediate as: round(expansion * hidden_size * 2/3)
                      then rounds up to nearest multiple of 256
        """
        super().__init__()
        # Calculate intermediate dimension with 2/3 scaling factor
        # Round to nearest multiple of 256 for hardware efficiency
        inter = _find_multiple(round(expansion * hidden_size * 2 / 3), 256)

        # Fused projection: projects to 2*inter, then chunks into gate and up
        # More memory efficient than two separate projections
        self.gate_up_proj = CastedLinear(hidden_size, inter * 2, bias=False)
        self.down_proj = CastedLinear(inter, hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input tensor of shape [..., hidden_size]

        Returns:
            Activated tensor of same shape as input
        """
        # Single fused projection, then split
        gate, up = self.gate_up_proj(x).chunk(2, dim=-1)
        # Apply SiLU to gate, multiply with up, project down
        return self.down_proj(F.silu(gate) * up)


class MLPMixer(nn.Module):
    """MLP-Mixer for token/sequence dimension mixing.

    Implements token mixing via transpose + SwiGLU pattern:
    - Transpose: [B,L,H] -> [B,H,L] to make sequence dimension accessible
    - Apply SwiGLU on sequence dimension (treats L as feature dim)
    - Transpose back: [B,H,L] -> [B,L,H]

    This allows information mixing across the sequence/token dimension,
    matching the MLP-Mixer architecture pattern from:
    "MLP-Mixer: An all-MLP Architecture for Vision" (Tolstikhin et al., 2021)

    Args:
        seq_len: Sequence length (dimension to mix across)
        expansion: Expansion factor for SwiGLU (default: 4.0)
    """

    def __init__(self, seq_len: int, expansion: float = 4.0):
        super().__init__()
        self.swiglu = SwiGLU(hidden_size=seq_len, expansion=expansion)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply token mixing.

        Args:
            x: [B, L, H] input tensor

        Returns:
            [B, L, H] output with token mixing applied
        """
        x = x.transpose(1, 2)  # [B,L,H] -> [B,H,L]
        x = self.swiglu(x)     # Apply SwiGLU on L dimension
        x = x.transpose(1, 2)  # [B,H,L] -> [B,L,H]
        return x


class NetworkBlock(nn.Module):
    """
    2-layer network block for TRM.

    As specified in Section 4.4 and Table 1 (TRM uses 2 layers, not 4 like HRM).

    Each layer consists of:
    1. RMSNorm
    2. MLPMixer (sequence mixing on length dimension)
    3. RMSNorm
    4. SwiGLU MLP (channel mixing)

    All hyperparameters hardcoded as specified in paper:
    - num_layers = 2
    - hidden_size = 512
    - context_length = 82 (1 summary token + 81 Sudoku tokens)

    No bias anywhere (Section 2.1).
    """

    def __init__(self):
        """
        Initialize 2-layer network block with hardcoded hyperparameters.
        """
        super().__init__()

        # Hardcoded hyperparameters from Table 1, Page 11
        self.hidden_size = 512
        self.num_layers = 2
        self.seq_len = 1 + 81  # 1 summary token + 81 Sudoku tokens
        self.norm_eps = 1e-5  # Epsilon for rms_norm (matching reference)

        # Create layers
        self.layers = nn.ModuleList([
            self._create_layer() for _ in range(self.num_layers)
        ])

    def _create_layer(self):
        """
        Create a single layer with token and channel mixing.

        Following MLP-Mixer architecture pattern:
        - mixer: MLPMixer for token mixing (operates on sequence dimension L=82)
        - mlp: SwiGLU for channel mixing (operates on hidden dimension H=512)
        """
        return nn.ModuleDict({
            'mixer': MLPMixer(seq_len=self.seq_len, expansion=4.0),
            'mlp': SwiGLU(hidden_size=self.hidden_size, expansion=4.0),
        })

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through 2-layer network with post-norm pattern.

        Follows MLP-Mixer architecture:
        - Token mixing: MLPMixer mixes information across sequence dimension
        - Channel mixing: SwiGLU mixes information across hidden dimension
        Both use post-norm (norm after residual connection).

        Args:
            x: Input tensor of shape [batch, seq_len, hidden_size]

        Returns:
            Output tensor of same shape as input
        """
        for layer in self.layers:
            # Token mixing
            out = layer['mixer'](x)
            x = rms_norm(x + out, eps=self.norm_eps)

            # Channel mixing
            out = layer['mlp'](x)
            x = rms_norm(x + out, eps=self.norm_eps)

        return x


class InputEmbedding(nn.Module):
    """
    Input embedding for TRM.

    Maps input tokens to hidden dimension and prepends a summary token.
    NO position embeddings - position information is learned implicitly through MLPMixer.

    Summary token design (matching reference implementation):
    - Single learnable 512-dim vector (zero-initialized)
    - Prepended to sequence: [summary, token_0, ..., token_80]
    - Used by Q-head for halting decisions

    All hyperparameters hardcoded:
    - hidden_size = 512
    - context_length = 81 (Sudoku tokens only, summary token added separately)
    """

    def __init__(self, vocab_size, contenxt_length):
        """Initialize input embedding with hardcoded hyperparameters.
        :param vocab_size:
        """
        super().__init__()

        # Hardcoded hyperparameters from Page 11
        self.vocab_size = vocab_size
        self.hidden_size = 512
        self.context_length = contenxt_length

        self.embed_scale = math.sqrt(self.hidden_size)

        # Token embedding with HRM-style scaled initialization
        # std = 1 / sqrt(hidden_size) for numerical stability (matches HRM implementation)
        self.token_embedding = nn.Embedding(self.vocab_size, self.hidden_size)
        embed_init_std = 1.0 / self.embed_scale  # ≈ 0.044
        trunc_normal_init_(self.token_embedding.weight, std=embed_init_std)

        # Summary token: zero-initialized learnable vector [1, 1, 512]
        # Matches reference implementation (init_std=0)
        # This token will aggregate global information during recursion
        self.summary_token = nn.Parameter(torch.zeros(1, 1, self.hidden_size))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Embed input tokens and prepend summary token.

        Args:
            x: Input tokens of shape [batch, 81] with values in [0, 10]

        Returns:
            Embedded tensor of shape [batch, 82, 512]
            where position 0 is summary token, positions 1-81 are Sudoku tokens
        """
        batch_size = x.shape[0]

        # Embed tokens: [batch, 81] -> [batch, 81, 512]
        token_emb = self.token_embedding(x)

        # Expand summary token to batch size: [1, 1, 512] -> [batch, 1, 512]
        summary = self.summary_token.expand(batch_size, -1, -1)

        # Prepend summary token: [batch, 1, 512] + [batch, 81, 512] -> [batch, 82, 512]
        embedding = torch.cat([summary, token_emb], dim=1)

        # Scale embeddings by sqrt(hidden_size) for numerical stability (HRM-style)
        # This balances embedding magnitude with initial states and prevents underflow
        return self.embed_scale * embedding


class OutputHead(nn.Module):
    """
    Output head for TRM.

    Maps hidden representations to vocabulary logits.

    Hardcoded parameters:
    - hidden_size = 512
    """

    def __init__(self, vocab_size):
        """Initialize output head with hardcoded hyperparameters.
        :param vocab_size:
        """
        super().__init__()

        # Hardcoded hyperparameters
        self.hidden_size = 512
        self.vocab_size = vocab_size

        # No bias as specified in paper Section 2.1
        # CastedLinear automatically applies LeCun normal initialization
        self.proj = CastedLinear(self.hidden_size, self.vocab_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Project hidden states to vocabulary logits.

        Args:
            x: Hidden states of shape [batch, seq_len, hidden_size]

        Returns:
            Logits of shape [batch, seq_len, vocab_size]
        """
        return self.proj(x)


class QHead(nn.Module):
    """
    Q head for halting probability prediction (ACT).

    As specified in Section 4.6 and Algorithm 3.
    Predicts whether the model should halt (solution is correct).

    Design (matching reference implementation):
    - Takes summary token vector as input (not full sequence)
    - Summary token is extracted by TRM before passing to Q-head
    - No pooling needed - operates on single vector per sample

    Hardcoded parameters:
    - hidden_size = 512
    - output_dim = 1 (binary halting decision)
    """

    def __init__(self):
        """Initialize Q head with hardcoded hyperparameters."""
        super().__init__()

        # Hardcoded hyperparameters
        self.hidden_size = 512

        # Q head with bias (special case, differs from paper's general "no bias" rule)
        # HRM reference implementation uses bias for Q head with special initialization
        self.proj = nn.Linear(self.hidden_size, 1, bias=True)

        # Special initialization for faster learning during bootstrapping (HRM-style)
        # Init weights to zero and bias to -5 (predicts low halting probability initially)
        with torch.no_grad():
            self.proj.weight.zero_()
            self.proj.bias.fill_(-5.0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Predict halting probability from summary token vector.

        Args:
            x: Summary token hidden state of shape [batch, hidden_size]
               (NOT the full sequence - summary token pre-extracted by TRM)

        Returns:
            Halting logits of shape [batch, 1]
        """
        # Project summary token to single logit (no pooling needed)
        return self.proj(x)  # [batch, 1]


class TRM(nn.Module):
    """
    Tiny Recursive Model for Sudoku.

    Complete TRM implementation following Algorithm 3 (Figure 3, Page 5) of:
    "Less is More: Recursive Reasoning with Tiny Networks" (arXiv:2510.04871v1)

    Architecture:
    - Input embedding: f_I maps tokens to hidden states
    - Single 2-layer network: shared for both z and y updates
    - Output head: f_O maps hidden states to vocabulary logits
    - Q head: predicts halting probability

    Recursive process:
    - Latent recursion: z ← net(x, y, z) repeated n=6 times, then y ← net(y, z)
    - Deep recursion: T=3 cycles (2 without gradients, 1 with gradients)

    All hyperparameters hardcoded (Table 1, Page 11):
    - hidden_size = 512
    - num_layers = 2
    - n_recursions = 6
    - T_cycles = 3
    - context_length = 81
    - vocab_size = 11
    """

    def __init__(
            self,
            vocab_size=11,  # tokens 0-10 (PAD + 0-9)
            context_length=81  # 9x9 Sudoku grid (excludes summary token)
    ):
        """Initialize TRM model with hardcoded hyperparameters."""
        super().__init__()
        self.context_length = context_length
        self.vocab_size = vocab_size

        # Hardcoded hyperparameters from Table 1, Page 11
        self.hidden_size = 512
        self.n = 6  # latent recursions per cycle
        self.T = 3  # deep recursion cycles

        # Sequence length including summary token
        self.seq_len = context_length + 1  # 82 = 1 summary + 81 Sudoku

        # Initialize embeddings, network, heads
        self.input_embedding = InputEmbedding(self.vocab_size, self.context_length)
        self.network = NetworkBlock()
        self.output_head = OutputHead(self.vocab_size)
        self.q_head = QHead()

        # Position-agnostic initial states [512] (matching reference implementation)
        # Single 1D vector broadcast to all sequence positions
        # Initialized with truncated normal (std=1.0) following reference implementation
        # This reduces parameters from 82K to 1K while enforcing position invariance
        y_init_tensor = torch.empty(self.hidden_size)
        z_init_tensor = torch.empty(self.hidden_size)
        trunc_normal_init_(y_init_tensor, std=1.0)
        trunc_normal_init_(z_init_tensor, std=1.0)
        self.register_buffer('y_init', y_init_tensor, persistent=True)
        self.register_buffer('z_init', z_init_tensor, persistent=True)

    def latent_recursion(
            self,
            x: torch.Tensor,
            y: torch.Tensor,
            z: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Perform n=6 latent reasoning steps, then update answer.

        Following Algorithm 3:
        - For i in range(n): z = net(x + y + z)  # x present: update z
        - y = net(y + z)                          # x absent: update y

        The single network distinguishes between tasks by the input it receives.

        Args:
            x: Input embeddings [batch, 82, 512] (includes summary token)
            y: Current answer state [batch, 82, 512]
            z: Latent reasoning state [batch, 82, 512]

        Returns:
            Tuple of (updated_y, updated_z), both [batch, 82, 512]
        """
        # Latent reasoning: update z n=6 times with x present
        for _ in range(self.n):
            # Sum x, y, z element-wise (matching paper's ⊕ symbol)
            z_input = x + y + z  # [batch, 82, 512]
            # Update z through network
            z = self.network(z_input)  # [batch, 82, 512]

        # Answer update: update y once with x absent
        # Sum y, z element-wise (no x)
        y_input = y + z  # [batch, 82, 512]
        # Update y through network
        y = self.network(y_input)  # [batch, 82, 512]

        return y, z

    def deep_recursion(
            self,
            x: torch.Tensor,
            y: torch.Tensor,
            z: torch.Tensor
    ) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]:
        """
        Perform T=3 deep recursion cycles.

        Following Algorithm 3:
        - T-1=2 times without gradients
        - 1 time with gradients
        - Detach states after final cycle
        - Extract summary token and Sudoku states
        - Get predictions from extracted states

        Args:
            x: Input embeddings [batch, 82, 512] (includes summary token)
            y: Initial answer state [batch, 82, 512]
            z: Initial reasoning state [batch, 82, 512]

        Returns:
            Tuple of:
            - (y_detached, z_detached): Detached states for next step [batch, 82, 512] each
            - y_hat: Prediction logits [batch, 81, vocab_size] (Sudoku positions only)
            - q_hat: Halting logit [batch, 1] (from summary token)
        """
        # T-1=2 cycles without gradients
        with torch.no_grad():
            for _ in range(self.T - 1):
                y, z = self.latent_recursion(x, y, z)

        # Final cycle with gradients
        y, z = self.latent_recursion(x, y, z)

        # Detach states to prevent gradient accumulation across supervision steps
        y_detached = y.detach()
        z_detached = z.detach()

        # Extract summary token and Sudoku states for predictions
        # Position 0: summary token (for Q-head)
        # Positions 1-81: Sudoku tokens (for output head)
        summary = y[:, 0]        # [batch, 512]
        sudoku_states = y[:, 1:]  # [batch, 81, 512]

        # Get predictions from extracted states (before detachment for gradients)
        y_hat = self.output_head(sudoku_states)  # [batch, 81, vocab_size]
        q_hat = self.q_head(summary)             # [batch, 1]

        return (y_detached, z_detached), y_hat, q_hat

    def forward(
        self,
        x_input: torch.Tensor,
        y: torch.Tensor | None = None,
        z: torch.Tensor | None = None,
    ) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]:
        """
        Main forward pass for one supervision step.

        Following Algorithm 3:
        1. Embed input tokens (adds summary token)
        2. Initialize or reuse states
        3. Perform deep recursion
        4. Extract summary token and Sudoku states
        5. Return states and predictions

        Supports two modes:
        - **Inference mode** (y=None, z=None): Initializes fresh states from learnable parameters
        - **Training mode** (y, z provided): Reuses states from previous supervision step

        Args:
            x_input: Input token indices [batch, 81] with values in [0, 9]
            y: Optional pre-initialized state [batch, 82, 512]. If None, uses learnable y_init
            z: Optional pre-initialized state [batch, 82, 512]. If None, uses learnable z_init

        Returns:
            Tuple of:
            - (y, z): Detached states for next supervision step [batch, 82, 512] each
            - y_hat: Prediction logits [batch, 81, vocab_size] (Sudoku positions only)
            - q_hat: Halting logit [batch, 1] (from summary token)
        """
        batch_size = x_input.size(0)

        # Embed input tokens (adds summary token at position 0)
        x = self.input_embedding(x_input)  # [batch, 82, 512]

        # Initialize states from position-agnostic parameters or reuse provided states
        if y is None:
            # Broadcast [512] to [batch, 82, 512] - all positions get same init vector
            y = self.y_init.view(1, 1, -1).expand(batch_size, self.seq_len, -1)
        if z is None:
            # Broadcast [512] to [batch, 82, 512] - all positions get same init vector
            z = self.z_init.view(1, 1, -1).expand(batch_size, self.seq_len, -1)

        # Perform deep recursion
        (y, z), y_hat, q_hat = self.deep_recursion(x, y, z)

        return (y, z), y_hat, q_hat