transformer.py

from dataclasses import dataclass, field
from math import inf
from typing import Optional, Tuple, Union, List
from copy import deepcopy
import torch
from torch import nn
from torch.nn import Parameter
from torch.nn import functional as F
from torch.nn.attention.flex_attention import (
    create_block_mask, 
    BlockMask,
    _mask_mod_signature, 
    noop_mask,
)

torch._dynamo.config.capture_scalar_outputs = True 

from .lingua_transformer import (
    RMSNorm,
    InitStdFactor,
    RotaryEmbedding,
    TransformerBlock,
)
from .xattn import DecoderBlock, FourierConditioner, DecoderArgs, AdaRMSNorm
from .conv_stem import CausalConv2DStem
from .bottlenecks import mmd_imq
from vector_quantize_pytorch import SimVQ, FSQ
import functools

# def create_causal_mask(seqlen, attn_impl, sliding_window):
#     if attn_impl == "sdpa":
#         return "causal"
#     elif attn_impl == "flex_attention":
#         return create_block_mask(causal_mask, None, None, seqlen, seqlen)
#     else:
#         raise NotImplementedError(
#             f"Attention {attn_impl} with {sliding_window} sliding window not implemented"
#         )


def create_document_mask(lengths: torch.Tensor,
                         kv_lengths: Optional[torch.Tensor] = None, # for cross-attn
                         base_mask_mod: Optional[_mask_mod_signature] = None):
    """
    Create a document mask. Grabbing code from lingua.transformer
    """

    def generate_doc_mask_mod(
        mask_mod: _mask_mod_signature,
        lengths: torch.Tensor,
        kv_lengths: Optional[torch.Tensor] = None, # for cross-attn
    ) -> _mask_mod_signature:
        """Generates mask mods that apply to inputs to flex attention in the sequence stacked
        format.

        Args:
            mask_mod: The mask mod to apply to the documents
            lengths: Lengths of each document

        Note:
            What is the sequence stacked format? When assembling batches of inputs, we
            take multiple sequences and stack them together to form 1 large sequence. We then
            use masking to ensure that the attention scores are only applied to tokens within
            the same document.

        Example:

        - Square mask
        doc_mask         lengths
        a a b b b c c    2 3 2
        a 1 0 0 0 0 0 0
        a 1 1 0 0 0 0 0
        b 0 0 1 0 0 0 0
        b 0 0 1 1 0 0 0
        b 0 0 1 1 1 0 0
        c 0 0 0 0 0 1 0
        c 0 0 0 0 0 1 1

        """

        def lengths_to_start_ids(lengths):
            doc_start = lengths.cumsum(0)
            doc_start = doc_start.roll(1)
            doc_start[0] = 0
            return doc_start

        def lengths_to_local_ids(lengths):
            assert lengths.ndim == 1
            nb_seqs = lengths.size(0)
            total_seqlen = lengths.sum()
            # This gives the document id of each token
            doc_id = torch.repeat_interleave(lengths)
            # Compute document start for each document
            doc_start = lengths_to_start_ids(lengths)
            # Compute document start for each token
            doc_start = doc_start[doc_id]
            # Compute the position of each token within each document
            tok_id = torch.arange(total_seqlen, device=lengths.device) - doc_start
            return doc_id, tok_id

        kv_lengths = kv_lengths if kv_lengths is not None else lengths
        q_document_id, q_token_id = lengths_to_local_ids(lengths)
        kv_document_id, kv_token_id = lengths_to_local_ids(kv_lengths)
        q_max_idx = lengths.sum() - 1
        kv_max_idx = kv_lengths.sum() - 1

        def doc_mask_mod(b, h, q_idx, kv_idx):        
            q_idx_cap = torch.minimum(q_max_idx, q_idx)
            kv_idx_cap = torch.minimum(kv_max_idx, kv_idx)
            valid_idx = (q_idx <= q_max_idx) & (kv_idx <= kv_max_idx)
            same_doc = q_document_id[q_idx_cap] == kv_document_id[kv_idx_cap]
            q_logical = q_token_id[q_idx_cap]
            kv_logical = kv_token_id[kv_idx_cap]
            inner_mask = mask_mod(b, h, q_logical, kv_logical)
            return same_doc & inner_mask & valid_idx

        return doc_mask_mod

    if base_mask_mod is None:
        base_mask_mod = noop_mask

    if torch.cuda.is_available():
        doc_mask_mod = generate_doc_mask_mod(base_mask_mod, lengths)
        return create_block_mask(doc_mask_mod, None, None, lengths.sum().item(), lengths.sum().item())
    else:
        # create_block_mask runs on CPU; ensure closure tensors are on CPU too
        doc_mask_mod = generate_doc_mask_mod(base_mask_mod, lengths.cpu())
        return create_block_mask(doc_mask_mod, None, None, lengths.sum().item(), lengths.sum().item(),
                                device='cpu', _compile=False)


def attention_flops_per_token(n_layers, seq_len, dim, causal):
    # Formula from https://github.com/Dao-AILab/flash-attention/blob/main/benchmarks/benchmark_flash_attention.py#L27-L30
    return 3.5 * (4 * n_layers * seq_len * dim // (2 if causal else 1))


def get_num_flop_per_token(
    num_non_embed_params: int, n_layers: int, dim: int, seq_len: int
) -> int:
    return 6 * num_non_embed_params + attention_flops_per_token(
        n_layers, seq_len, dim, True
    )


def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx


def extract_non_registers(h: torch.Tensor, num_groups: int, original_seqlen: int = None, downsample_factor: int = None) -> torch.Tensor:
    """
    Extract non-register tokens from the output tensor.
    Args:
        h: Output tensor from transformer layers [B, interleaved_seqlen, D].
        num_groups: Number of groups used in interleaving.
        original_seqlen: The sequence length of the input *before* padding
                            and interleaving. Used to trim non-register tokens.

    Returns:
        non_registers: [B, original_seqlen, D]
    """
    bsz, seq_len, dim = h.shape
    # seq_len should be num_groups*(df+1)
    h = h.reshape(bsz, num_groups, downsample_factor + 1, dim)

    # Extract non-register tokens (indices 1 to df)
    non_registers = h[:, :, 1:, :]
    # Flatten back to sequence dimension
    padded_seqlen = num_groups * downsample_factor
    non_registers = non_registers.reshape(bsz, padded_seqlen, dim)
    # Trim back to the original sequence length, removing padding effects
    non_registers = non_registers[:, :original_seqlen, :]  # [B, original_seqlen, D]

    return non_registers.contiguous()

@torch.compile()
def huber_loss(target, logits, huber_c):
    return huber_c * (torch.sqrt((target - logits) ** 2 + huber_c**2) - huber_c)

@torch.compile()
def cosine_similarity_loss(input, target):
    return (1 - F.cosine_similarity(input, target, dim=-1).mean())

@torch.compile()
def huber_cosine_weighted(input, target, huber_c = 0.1):
    # Compute the Huber loss
    h_loss = huber_loss(input, target, huber_c).mean() * 0.5
    
    # Compute the cosine similarity
    cosine_sim = cosine_similarity_loss(input, target)
    
    # Combine the two losses
    combined_loss = h_loss + cosine_sim
    
    return combined_loss

@dataclass
class DecoderTransformerArgs(DecoderArgs):

    seed: int = 42

    weight_tying: bool = False
    sliding_window: int = 128
    xattn_sliding_window: int = 32
    input_dim: int = 64 

    decoder_encoder_dropout: float = 0.1
    decoder_timestep_dropout: float = 0.1

    encoder_sliding_window: int = 128
    encoder_input_dim: int = input_dim 
    encoder_output_dim: int = input_dim*2
    encoder_latent_downsample_factor: int = 2  
    encoder_hidden_dim: Optional[int] = None

    adaptive_loss_weighting: bool = False  
    num_fine_time_pts: int = 128
    dont_noise_chan_xyz: bool = False
    stft_global_sigma: Union[str, float] = 1.0 

    dropout_type: str =  "zero" # {"zero", "rand", "learnable"}

    bottleneck_type: str = "mmd"
    distill_output_dim: int = 0
    codebook_size: int = 1024
    levels: List[int] = field(default_factory=list)
    init_base_std: float = 0.02
    learnable_bias: bool = False

    huber_c: Optional[float] = None

    decoder_repa_index: float = float('inf')
    encoder_repa_index: float = float('inf')
    repa_dim: int = 1024
    repa_loss_fn: str = "cosine"

    compression_free_conv_stem: bool = False

    max_seqlen: int = 1024

class BaseTransformerDecoder(nn.Module):
    def __init__(self, args: DecoderTransformerArgs):
        super().__init__()

        self.dim = args.dim
        self.init_base_std = args.init_base_std
        self.init_std_factor = InitStdFactor(args.init_std_factor)
        self.max_seqlen = args.max_seqlen
        self.rope_embeddings = RotaryEmbedding(
            theta=args.rope_theta,
            head_dim=args.head_dim or args.dim // args.n_heads,
            max_seqlen=args.max_seqlen,
            rope_dim=args.rope_dim,
        )

        self.layers = nn.ModuleList()
        for _ in range(args.n_layers):
            self.layers.append(DecoderBlock(args))

        self.repa_index = args.decoder_repa_index
        if self.repa_index != inf:
            # self.repa_proj = nn.Linear(args.dim, args.repa_dim)
            self.repa_proj = nn.Sequential(
                nn.Linear(args.dim, args.repa_dim),
                nn.SiLU(),
                nn.Linear(args.repa_dim, args.repa_dim),
                nn.SiLU(),
                nn.Linear(args.repa_dim, args.repa_dim),
            )
            self.repa_norm = AdaRMSNorm(args.t_dim, args.dim, eps=args.norm_eps)
            self.repa_loss_fn = cosine_similarity_loss if args.repa_loss_fn == "cosine" else huber_cosine_weighted

    def forward(
        self,
        h,
        x_attended,
        t,
        tok_idx: Optional[torch.Tensor] = None,
        cross_tok_idx: Optional[torch.Tensor] = None,
        mask: Optional[Union[BlockMask,  str]] = None,
        cross_attn_mask: Optional[Union[BlockMask,  str]] = None,
        attn_impl: str = "sdpa",
        repa_target: Optional[torch.Tensor] = None,
        do_idx: Optional[torch.Tensor] = None,
    ):

        freq_cis = self.rope_embeddings(seqlen=self.max_seqlen, tok_idx=tok_idx)
        repa_loss = None

        for i, layer in enumerate(self.layers):     # all these layers are type 'xattn.DecoderBlock'
            h = layer(h,
                      x_attended,
                      t,
                      freq_cis,
                      tok_idx=tok_idx,
                      cross_tok_idx=cross_tok_idx,
                      self_attn_mask=mask,
                      cross_attn_mask=cross_attn_mask,
                      attn_impl=attn_impl,
                      do_idx=do_idx,
            )

            if self.training and self.repa_index != inf and i == self.repa_index:
                repa_loss = self.repa_loss_fn(self.repa_proj(self.repa_norm(h, t)).float(), repa_target,)

        return h, repa_loss



    def reset_parameters(self):
        # Either use fixed base std or sqrt model dim
        self.rope_embeddings.reset_parameters()

    def init_weights(self):
        self.reset_parameters()
        for depth, layer in enumerate(self.layers):
            factor = {
                InitStdFactor.CURRENT_DEPTH: (2 * (depth + 1)) ** 0.5,
                InitStdFactor.GLOBAL_DEPTH: (2 * (len(self.layers) + 1)) ** 0.5,
                InitStdFactor.DIM_RATIO: self.dim / 4096,
                InitStdFactor.DISABLED: 1.0,
            }[self.init_std_factor]

            layer.init_weights(self.init_base_std, factor)

        # Add these lines for repa_proj initialization
        if self.repa_index != float('inf'):
            init_std = self.init_base_std or (self.dim ** (-0.5))
            self.repa_norm.reset_parameters() # Ensure repa_norm is also reset
            
            #now repa_proj is nn.Sequential, let's do it in a loop
            for i, layer in enumerate(self.repa_proj):
                if isinstance(layer, nn.Linear):
                    nn.init.trunc_normal_(
                        layer.weight,
                        mean=0.0,
                        std=init_std,
                        a=-3 * init_std,
                        b=3 * init_std,
                    )
                    # nn.init.zeros_(layer.weight)
                    if layer.bias is not None:
                        nn.init.zeros_(layer.bias)
            

class BaseTransformer(nn.Module):
    def __init__(self, args: DecoderTransformerArgs):
        super().__init__()

        self.dim = args.dim
        self.init_base_std = args.init_base_std
        self.init_std_factor = InitStdFactor(args.init_std_factor)
        self.max_seqlen = args.max_seqlen
        self.rope_embeddings = RotaryEmbedding(
            theta=args.rope_theta,
            head_dim=args.head_dim or args.dim // args.n_heads,
            max_seqlen=args.max_seqlen,
            rope_dim=args.rope_dim,
        )

        self.layers = nn.ModuleList()
        for _ in range(args.n_layers):
            self.layers.append(TransformerBlock(args))

        self.repa_index = args.encoder_repa_index
        if self.repa_index != inf:
            self.repa_proj = nn.Linear(args.dim, args.repa_dim)
            self.repa_norm = RMSNorm(args.dim, eps=args.norm_eps)
            self.repa_loss_fn = cosine_similarity_loss if args.repa_loss_fn == "cosine" else huber_cosine_weighted


    def forward(
        self,
        h,
        tok_idx: Optional[torch.Tensor] = None,
        mask: Optional[Union[BlockMask,  str]] = None,
        attn_impl: str = "sdpa",
        repa_target: Optional[torch.Tensor] = None,
        do_idx: Optional[torch.Tensor] = None,
        **kwargs
    ):

        freq_cis = self.rope_embeddings(seqlen=self.max_seqlen,
                                        tok_idx=tok_idx
        )
        repa_loss = None

        for i, layer in enumerate(self.layers):     # all these layers are type 'TransformerBlock'
            h = layer(h, 
                      freq_cis, 
                      tok_idx=tok_idx, 
                      mask=mask, 
                      attn_impl=attn_impl,
                      do_idx=do_idx,
            )

            if self.training and self.repa_index != inf and i == self.repa_index:
                repa_loss = cosine_similarity_loss(self.repa_proj(self.repa_norm(extract_non_registers(h, **kwargs))).float(), repa_target,)


        return h, repa_loss

    def reset_parameters(self):
        # Either use fixed base std or sqrt model dim
        self.rope_embeddings.reset_parameters()

    def init_weights(self):
        self.reset_parameters()
        for depth, layer in enumerate(self.layers):
            factor = {
                InitStdFactor.CURRENT_DEPTH: (2 * (depth + 1)) ** 0.5,
                InitStdFactor.GLOBAL_DEPTH: (2 * (len(self.layers) + 1)) ** 0.5,
                InitStdFactor.DIM_RATIO: self.dim / 4096,
                InitStdFactor.DISABLED: 1.0,
            }[self.init_std_factor]

            layer.init_weights(self.init_base_std, factor)

        # Add these lines for repa_proj initialization
        if self.repa_index != float('inf'):
             init_std = self.init_base_std or (self.dim ** (-0.5))

             nn.init.trunc_normal_(
                 self.repa_proj.weight,
                 mean=0.0,
                 std=init_std, 
                 a=-3 * init_std,
                 b=3 * init_std,
             )
             if self.repa_proj.bias is not None:
                 nn.init.zeros_(self.repa_proj.bias)
             self.repa_norm.reset_parameters() 

class DecoderTransformer(BaseTransformerDecoder):
    def __init__(self, args: DecoderTransformerArgs):
        super().__init__(args)
        self.weight_tying = False
        self.sliding_window = args.sliding_window
        self.xattn_sliding_window = args.xattn_sliding_window
        if args.huber_c is not None:
            self.huber_c = args.huber_c
        else:
            self.huber_c = None 

        self.tok_embeddings = nn.Linear(args.input_dim, args.dim,) 

        self.t_embedder = FourierConditioner(args.t_dim, std=args.init_base_std)

        self.encoder_proj = nn.Linear(args.encoder_output_dim, args.dim)

        self.norm = AdaRMSNorm(args.t_dim, args.dim, eps=args.norm_eps)

        self.output = nn.Linear(
            args.dim,
            args.input_dim, 
            bias=False,
        )
        self.init_base_std = args.init_base_std
        self.use_compression_free_conv_stem = False
        if args.compression_free_conv_stem:
            self.use_compression_free_conv_stem = True

            self.compression_free_conv_stem_input = CausalConv2DStem(
                input_features = args.input_dim, 
                hidden_channels = 32,
                activation = nn.SELU,
                compress_channels=False,
            )
            self.compression_free_conv_stem_output = CausalConv2DStem(
                input_features = args.input_dim,
                hidden_channels = 32,
                activation = nn.SELU,
                compress_channels=False,
            )

        self.adaptive_loss_weighting = args.adaptive_loss_weighting


    def forward(
        self,
        tokens: torch.Tensor,
        cross_attended: torch.Tensor,
        timeD: torch.Tensor,
        seq_lens: torch.Tensor,         # for document masking packed sequences in self-attention
        cross_seq_lens: torch.Tensor,   # for document masking packed sequences in cross-attention
        target: Optional[torch.Tensor] = None,
        tok_idx: Optional[torch.Tensor] = None,
        cross_tok_idx: Optional[torch.Tensor] = None,
        mask: Optional[Union[BlockMask,  torch.Tensor, str]] = None,
        cross_attn_mask: Optional[Union[BlockMask,  str]] = None,
        attn_impl: str = "flex_attention",
        time_masks: Optional[torch.Tensor] = None,
        channel_loss_weighting: Optional[torch.Tensor] = None, # [1, 1, input_dim*2]
        repa_target: Optional[torch.Tensor] = None,
        freq_masks: Optional[torch.Tensor] = None,
        do_idx: Optional[torch.Tensor] = None,
        print_layerwise_activation_stats: bool = False,
    ):

        tokens = tokens.squeeze(1)  

        bsz, seqlen, dim = tokens.shape
        _, cross_seqlen, _ = cross_attended.shape

        # Masking out channels that were set to all-zeros
        if self.training and freq_masks is not None:
            with torch.no_grad():
                tokens *= freq_masks
        
        if self.use_compression_free_conv_stem:
            tokens = self.compression_free_conv_stem_input(tokens)

        h = self.tok_embeddings(tokens)
        t = self.t_embedder(timeD)

        cross_attended = self.encoder_proj(cross_attended)


        # COMBINE SLIDING WINDOW MASK AND DOCUMENT MASK
        SLIDING_WINDOW = self.sliding_window
        def selfattn_sliding_window_func(b, h, q_idx, kv_idx):
            # Self-attention case
            return (q_idx - kv_idx).abs() <= SLIDING_WINDOW
        mask_mod_slide = selfattn_sliding_window_func
        mask = create_document_mask(lengths=seq_lens, base_mask_mod=mask_mod_slide)
        #
        SLIDING_WINDOW = self.xattn_sliding_window
        def crossattn_sliding_window_func(b, h, q_idx, kv_idx):
            # Cross-attention case
            center_k_idx = (q_idx * cross_seqlen) // seqlen
            return (kv_idx - center_k_idx).abs() <= SLIDING_WINDOW
        mask_mod_slide = crossattn_sliding_window_func
        cross_attn_mask = create_document_mask(lengths=seq_lens, kv_lengths=cross_seq_lens, base_mask_mod=mask_mod_slide)

        visualize_attention_masks = False
        if visualize_attention_masks:
            from .utils import visualize_attention_mask
            torch._dynamo.config.disable = True
            visualize_attention_mask(mask, title_suffix="decoder_self_attn")
            visualize_attention_mask(cross_attn_mask, title_suffix="decoder_cross_attn")
            torch._dynamo.config.disable = False


        if tok_idx is not None:
            if tok_idx.ndim==3 and tok_idx.shape[0]==1:
                tok_idx = tok_idx.squeeze().squeeze() # make it the right size for RoPE.

        if cross_tok_idx is not None:
            if cross_tok_idx.ndim==3 and cross_tok_idx.shape[0]==1:
                cross_tok_idx = cross_tok_idx.squeeze().squeeze() # make it the right size for RoPE.

        h, repa_loss = super().forward(h,
                                       cross_attended,
                                       t=t,
                                       tok_idx=tok_idx,
                                       cross_tok_idx=cross_tok_idx,
                                       mask=mask, 
                                       cross_attn_mask=cross_attn_mask, 
                                       attn_impl=attn_impl,
                                       repa_target=repa_target,
                                       do_idx=do_idx,
        )

        h_normed = self.norm(h, t) 

        # if print_layerwise_activation_stats and do_idx is not None:
        #     print(f"\nDecoder output norm: (drop-out) mean={h[:, do_idx, :].mean().item():.6f}, std={h[:, do_idx, :].std().item():.6f}", end=" --> ")
        #     print(f"mean={h_normed[:, do_idx, :].mean().item():.6f}, std={h_normed[:, do_idx, :].std().item():.6f}")
        #     print(f"Decoder output norm: (non-drop) mean={h[:, ~do_idx, :].mean().item():.6f}, std={h[:, ~do_idx, :].std().item():.6f}", end=" --> ")
        #     print(f"mean={h_normed[:, ~do_idx, :].mean().item():.6f}, std={h_normed[:, ~do_idx, :].std().item():.6f}")

        logits = self.output(h_normed) 

        if self.use_compression_free_conv_stem:
            logits = self.compression_free_conv_stem_output(logits)

        losses = self.compute_losses(target, logits, time_masks, freq_masks, channel_loss_weighting)

        if repa_target is not None:
            losses["decoder_repa_loss"] = repa_loss#.mean()

        return logits, losses


    
    @torch.compile()
    def compute_losses(self, target, logits, time_masks, freq_masks, channel_loss_weighting):
        losses = {}

        if target is not None:
            if self.huber_c is None:
                batchwise_loss = F.mse_loss(target.float(), logits.float(), reduction="none") # shape = [B, T, C]
            else:
                batchwise_loss = huber_loss(target.float(), logits.float(), self.huber_c)

            #  Do Adaptive Loss Weighting - to boost loss from channels with small signals so we can better learn small signals.
            if self.adaptive_loss_weighting:
                ALW = batchwise_loss.detach().abs().mean(dim=2).unsqueeze(2) # shape = [B,C,1]
                batchwise_loss = batchwise_loss/(ALW + 1e-5)


            if channel_loss_weighting is not None:
                batchwise_loss = batchwise_loss * channel_loss_weighting

            if freq_masks is not None:
                batchwise_loss = (batchwise_loss * freq_masks).sum(dim=1) / (freq_masks.sum(dim=1) + 1e-6) # shape = [B,C,1]
            else:
                batchwise_loss = batchwise_loss.mean(dim=-1)

            if time_masks is not None:
                batchwise_loss = (batchwise_loss * time_masks).sum(dim=-1) / time_masks.sum(dim=-1)

            losses["decoder_rf_loss"] = batchwise_loss.mean()

        return losses


    def reset_parameters(self, init_std=None):
        # Either use fixed base std or sqrt model dim
        super().reset_parameters()
        if init_std is None:
            init_std = self.init_base_std or (self.dim ** (-0.5))

        self.norm.reset_parameters()
        self.t_embedder.reset_parameters(init_std)
        nn.init.trunc_normal_(
            self.tok_embeddings.weight,
            mean=0.0,
            std=init_std,
            a=-3 * init_std,
            b=3 * init_std,
        )

        nn.init.trunc_normal_(
            self.encoder_proj.weight,
            mean=0.0,
            std=init_std,
            a=-3 * init_std,
            b=3 * init_std,
        )
        if self.encoder_proj.bias is not None:
            nn.init.zeros_(self.encoder_proj.bias)

        nn.init.trunc_normal_(
            self.output.weight,
            mean=0.0,
            std=init_std,
            a=-3 * init_std,
            b=3 * init_std,
        )
        if self.use_compression_free_conv_stem:
            self.compression_free_conv_stem_input.reset_parameters(init_std)
            self.compression_free_conv_stem_output.reset_parameters(init_std)
            
    def init_weights(self):
        super().init_weights()
        self.reset_parameters()

class EncoderTransformer(BaseTransformer):
    def __init__(self, args: DecoderTransformerArgs):
        args = deepcopy(args)
        args.dim = args.dim if args.encoder_hidden_dim is None else args.encoder_hidden_dim
        super().__init__(args)
        self.weight_tying = False
        self.sliding_window = args.encoder_sliding_window
        self.bottleneck_type = args.bottleneck_type
        self.downsample_factor = args.encoder_latent_downsample_factor
        self.distill = args.distill_output_dim != 0
        self.tok_embeddings = nn.Linear(args.encoder_input_dim, args.dim) 
        self.norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.registers = torch.nn.Parameter(torch.zeros(1, args.encoder_input_dim))
        self.dropout_type = args.dropout_type
        if self.dropout_type=="learnable":
            self.dropout_vec = torch.nn.Parameter(args.stft_global_sigma*torch.rand(1, args.encoder_input_dim, dtype=torch.float32)) # rand init for learnable dropout vector
        else:
            self.dropout_vec = None # If None, it will just use zeros for dropped out chans (rather than learnable vector).

        self.init_base_std = args.init_base_std
        self.output = nn.Linear(args.dim, args.encoder_output_dim, bias=False)

        if args.distill_output_dim != 0:
            self.distill_output = nn.Linear(
                args.dim,
                args.distill_output_dim,
                bias=True,
            )
            self.distill_norm = RMSNorm(args.dim, eps=args.norm_eps)

        if "sim" in args.bottleneck_type:
            self.quantizer = SimVQ(
                dim = args.encoder_output_dim,
                codebook_size = args.codebook_size,
                rotation_trick = True  # use rotation trick from Fifty et al.
            )
        elif "fsq" in args.bottleneck_type:
            self.quantizer = FSQ(
                levels = args.levels
            )

        self.use_compression_free_conv_stem = False
        if args.compression_free_conv_stem:
            self.use_compression_free_conv_stem = True

            self.compression_free_conv_stem_input = CausalConv2DStem(
                input_features = args.input_dim,
                hidden_channels = 32,
                activation = nn.SELU,
                compress_channels=False,
            )


    def _interleave_registers(self, x: torch.Tensor):
        """
        1) Pad `x` along the sequence dimension so it’s divisible by `self.downsample_factor`.
        2) Reshape into groups of length `df`.
        3) Insert a copy of `self.registers` in front of each group.
        4) Flatten back out.

        Returns:
            interleaved: [B, num_groups*(df+1), D]
            num_groups: int
        """

        bsz, seqlen, dim = x.shape
        df = self.downsample_factor

        # Number of groups
        num_groups = (seqlen + df - 1) // df
        new_seqlen = num_groups * df

        # Pad if needed
        if new_seqlen > seqlen:
            pad_len = new_seqlen - seqlen
            x = torch.cat([x, x.new_zeros(bsz, pad_len, dim)], dim=1)

        # Reshape to [B, num_groups, df, D]
        x = x.reshape(bsz, num_groups, df, dim)

        # Expand the single register => [B, num_groups, 1, D]
        regs = self.registers.expand(bsz, num_groups, -1).unsqueeze(2)

        # Cat the register in front of each group => [B, num_groups, df+1, D]
        x = torch.cat([regs, x], dim=2)

        # Flatten => [B, num_groups*(df+1), D]
        x = x.reshape(bsz, -1, dim).contiguous()
        return x, num_groups

    def _extract_registers_and_non_registers(
        self,
        h: torch.Tensor,
        num_groups: int,
        original_seqlen: int = None,                                # Added: original sequence length before padding
        return_non_registers: bool = False                          # Added: flag to return other tokens
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:    # Updated return type hint
        """
        Args:
            h: Output tensor from transformer layers [B, interleaved_seqlen, D].
            num_groups: Number of groups used in interleaving.
            original_seqlen: The sequence length of the input *before* padding
                             and interleaving. Used to trim non-register tokens.
            return_non_registers: If True, return both register and non-register
                                  tokens. Otherwise, return only register tokens.

        Returns:
            If return_non_registers is False:
                registers: [B, num_groups, D]
            If return_non_registers is True:
                registers: [B, num_groups, D]
                non_registers: [B, original_seqlen, D]
        """
        bsz, seq_len, dim = h.shape
        df = self.downsample_factor
        # seq_len should be num_groups*(df+1)
        h = h.reshape(bsz, num_groups, df + 1, dim)

        # The register is the first token in each group
        registers = h[:, :, 0, :]  # [B, num_groups, D]

        if not return_non_registers:
            return registers.contiguous(), None
        else:
            # Extract non-register tokens (indices 1 to df)
            non_registers = h[:, :, 1:, :] # [B, num_groups, df, D]
            # Flatten back to sequence dimension
            padded_seqlen = num_groups * df
            non_registers = non_registers.reshape(bsz, padded_seqlen, dim)
            # Trim back to the original sequence length, removing padding effects
            non_registers = non_registers[:, :original_seqlen, :] # [B, original_seqlen, D]

            return registers.contiguous(), non_registers.contiguous()
        


    def forward(
        self,
        token_values: torch.Tensor,
        seq_lens: torch.Tensor,
        distill_target: Optional[torch.Tensor] = None,
        tok_idx: Optional[torch.Tensor] = None,
        mask: Optional[Union[BlockMask,  torch.Tensor, str]] = None,
        attn_impl: str = "flex_attention",
        repa_target: Optional[torch.Tensor] = None,
        do_idx: Optional[torch.Tensor] = None,
        print_layerwise_activation_stats: bool = False,
    ): 
        
        _, orig_seqlen, _ = token_values.shape

        if self.use_compression_free_conv_stem:
            token_values = self.compression_free_conv_stem_input(token_values)

        token_values, num_groups = self._interleave_registers(token_values)
        bsz, seqlen, _ = token_values.shape

        if do_idx is not None: # (CW)
            do_idx_pre_reg = do_idx                             # indices of dropped-out channels without registers interleaved       
            do_idx = (token_values.sum(axis=2)==0).squeeze(0)   # recompute do_idx after interleaving registers


        # Now if using Learable Dropout, replace dropped-out channels with a learned but fixed parameter vector.
        if self.dropout_vec is not None:
            token_values[:,do_idx,:] = self.dropout_vec

        h = self.tok_embeddings(token_values)

        # (CW) - COMBINE SLIDING WINDOW MASK AND DOCUMENT MASK 
        SLIDING_WINDOW = self.sliding_window
        def sliding_window_func(b, h, q_idx, kv_idx):
            # Self-attention case
            return (q_idx - kv_idx).abs() <= SLIDING_WINDOW
        mask_mod_slide = sliding_window_func
        mask = create_document_mask(seq_lens*2, base_mask_mod=mask_mod_slide) # Hardcoding for CR=1 with interleave_registers thing.

        visualize_attention_masks = False
        if visualize_attention_masks:
            from .utils import visualize_attention_mask
            torch._dynamo.config.disable = True
            visualize_attention_mask(mask, title_suffix="encoder")
            torch._dynamo.config.disable = False
            
        if tok_idx is not None:
            tok_idx = tok_idx.repeat_interleave(repeats=2,dim=1)
            tok_idx = tok_idx.squeeze().squeeze() # make it the right size for RoPE.


        # if print_layerwise_activation_stats and do_idx is not None: # (CW)
        #     print(f"{do_idx.sum()=} and {(~do_idx).sum()=}")
        #     print(f"{token_values.shape=}")


        h, repa_loss = super().forward(h,                   # BaseTransformer.forward
                                       tok_idx=tok_idx, 
                                       mask=mask,
                                       attn_impl=attn_impl, 
                                       repa_target=repa_target, 
                                       num_groups=num_groups, 
                                       original_seqlen=orig_seqlen, 
                                       downsample_factor=self.downsample_factor,
                                       do_idx=do_idx,
        )

        h, non_regs = self._extract_registers_and_non_registers(h, num_groups, original_seqlen=orig_seqlen, return_non_registers=distill_target is not None)


        # if print_layerwise_activation_stats and do_idx is not None: # (CW)
        #     h_normed = self.norm(h) # (CW)
        #     print(f"\nEncoder output norm (drop-out): mean={h[:, do_idx_pre_reg, :].mean().item():.6f}, std={h[:, do_idx_pre_reg, :].std().item():.6f}", end=" --> ") # (CW)
        #     print(f"mean={h_normed[:, do_idx_pre_reg, :].mean().item():.6f}, std={h_normed[:, do_idx_pre_reg, :].std().item():.6f}") # (CW)
        #     print(f"Encoder output norm (non-drop): mean={h[:, ~do_idx_pre_reg, :].mean().item():.6f}, std={h[:, ~do_idx_pre_reg, :].std().item():.6f}", end=" --> ") # (CW)
        #     print(f"mean={h_normed[:, ~do_idx_pre_reg, :].mean().item():.6f}, std={h_normed[:, ~do_idx_pre_reg, :].std().item():.6f}") # (CW)
        #     logits = self.output(h_normed) # (CW)
        logits = self.output(self.norm(h))


        logits, losses = self.bottleneck(logits)
        if distill_target is not None:
            losses['encoder_distill'] = (1 - F.cosine_similarity(self.distill_output(self.distill_norm(non_regs)), distill_target, dim=-1).mean()) * 0.1 

        if repa_target is not None:
            losses["encoder_repa_loss"] = repa_loss#.mean()
        return logits, losses

    def bottleneck(self, h,):
        losses = {}
        latent = h
        b, l, d = h.shape
        if "kl" in self.bottleneck_type:
            mean, logvar = h.chunk(2, dim=-1)
            logvar = logvar.clamp(min=-3)
            std = torch.exp(0.5 * logvar)
            latent = mean + std * torch.randn_like(mean)
            kl_div_loss = -0.5 * (1 + logvar - mean.pow(2) - logvar.exp())
            losses["kl"] = kl_div_loss.mean()

        if "mmd" in self.bottleneck_type and self.training:
            losses["mmd"] = mmd_imq(latent.view(b*l, d).float(), torch.randn((b*l,d), dtype=torch.float32, device=latent.device,), 10.0) 

        if "sim" in self.bottleneck_type:
            latent, codes, simvq_loss = self.quantizer(h)
            losses["simvq_commit_loss"] = simvq_loss

        if "fsq" in self.bottleneck_type:
            latent, codes = self.quantizer(h)

        return latent, losses

    def reset_parameters(self, init_std=None):
        # Either use fixed base std or sqrt model dim
        super().reset_parameters()
        if init_std is None:
            init_std = self.init_base_std or (self.dim ** (-0.5))
        self.norm.reset_parameters()
        nn.init.trunc_normal_(
            self.tok_embeddings.weight,
            mean=0.0,
            std=init_std,
            a=-3 * init_std,
            b=3 * init_std,
        )

        nn.init.trunc_normal_(
            self.output.weight,
            mean=0.0,
            std=init_std,
            a=-3 * init_std,
            b=3 * init_std,
        )
        
        if self.distill:
            self.distill_norm.reset_parameters()
            nn.init.trunc_normal_(
                self.distill_output.weight,
                mean=0.0,
                std=init_std,
                a=-3 * init_std,
                b=3 * init_std,
            )
            nn.init.zeros_(self.distill_output.bias)

        nn.init.trunc_normal_(
            self.registers,
            mean=0.0,
            std=init_std,
            a=-3 * init_std,
            b=3 * init_std,
        )
        if self.use_compression_free_conv_stem:
            self.compression_free_conv_stem_input.reset_parameters(init_std)

    def init_weights(self):
        super().init_weights()
        self.reset_parameters()

class EncoderDecoder(nn.Module):
    def __init__(self, args: DecoderTransformerArgs):
        super().__init__()

        self.encoder = EncoderTransformer(args)
        self.decoder = DecoderTransformer(args)
        self.input_output_dim = args.input_dim
        self.encoder_dropout = args.decoder_encoder_dropout
        self.decoder_timestep_dropout = args.decoder_timestep_dropout
        self.global_sigma = args.stft_global_sigma

        self.num_fine_time_pts = args.num_fine_time_pts
        self.dont_noise_chan_xyz = args.dont_noise_chan_xyz
        self.rope_dim = args.rope_dim
        self.tok_idx_type = args.tok_idx_type



    def forward(
            self,
            encoder_input: torch.Tensor,
            decoder_input: torch.Tensor,
            t: torch.Tensor,
            chan_pos: torch.Tensor,
            chan_pos_discrete: torch.Tensor,
            chan_id: torch.Tensor,
            t_coarse: torch.Tensor,
            seq_lens: torch.Tensor,
            target: Optional[torch.Tensor] = None,
            distill_target: Optional[torch.Tensor] = None,
            time_masks: Optional[torch.Tensor] = None,
            channel_loss_weighting: Optional[torch.Tensor] = None, # [1, 1, input_dim*2]
            encoder_repa_target: Optional[torch.Tensor] = None,
            decoder_repa_target: Optional[torch.Tensor] = None,
            freq_masks: Optional[torch.Tensor] = None, # 0s where to not compute loss, 1s where to compute loss [B, 1, C]               
    ):


        if encoder_input.ndim==2:
            encoder_input = encoder_input.unsqueeze(0)
            target = target.unsqueeze(0) # doing to get rid of broadcast warning from DecoderTransformer.compute_losses
            chan_pos = chan_pos.unsqueeze(0)
            chan_pos_discrete = chan_pos_discrete.unsqueeze(0)
            chan_id = chan_id.unsqueeze(0)
            t_coarse = t_coarse.unsqueeze(0)


        ## Options for tok_idx.  Choose 1.
        if self.tok_idx_type is None:
            tok_idx = None          # this will just use args.model.max_seqlen to construct 1D-RoPE (but requires max_seqlen way too long).
        elif self.tok_idx_type == "t_coarse" and self.rope_dim==1:
            tok_idx = t_coarse      # this ignores channel and just uses coarse time in 1D-RoPE
        elif self.tok_idx_type == "chan_id" and self.rope_dim==1:
            tok_idx = chan_id       # this uses channel id in 1D-RoPE                             # this is same as hstack(arange(seq_lens)) below when seq_len = num_chans, ie chop_signals_only
        elif self.tok_idx_type == "stack_arange_seqlen" and self.rope_dim==1:
            tok_idx = torch.hstack([torch.arange(sl) for sl in seq_lens]).unsqueeze(0).unsqueeze(-1) # This has a different tok_id value for each element in sequence (chan or tc).
        elif self.tok_idx_type == "{x,y,z,tc}" and self.rope_dim==4: 
            tok_idx = torch.cat((chan_pos_discrete,t_coarse), dim=2)
        else:
            raise ValueError(f"Dont understand {self.tok_idx_type=} and {self.rope_dim}")


        do_idx = (encoder_input.sum(axis=2)==0).squeeze(0) # indices of dropped-out channels (CW) 
        # do_idx = None # [Set do_idx to None to disable printing of activation stats comparing channel drop-out]

        enc_out, enc_losses = self.encoder(encoder_input, 
                                           distill_target=distill_target,       # (CW) - None
                                           repa_target=encoder_repa_target,     # (CW) - None
                                           mask=None,
                                           seq_lens=seq_lens,                   # (CW) - for document masking
                                           tok_idx=tok_idx,                    # (CW) - pass in coarse time index for 1D RoPE
                                           do_idx=do_idx,                      # indices of dropped-out channels (CW)
        )

        dec_out, dec_losses = self.decoder(tokens=decoder_input,
                                           cross_attended=enc_out,
                                           timeD=t, 
                                           target=target, 
                                           time_masks=time_masks,                               # (CW) - None
                                           channel_loss_weighting=channel_loss_weighting,       # (CW) - None
                                           repa_target=decoder_repa_target,                     # (CW) - None
                                           freq_masks = freq_masks,                             # (CW) - masks out bad (all-zero) channels [B, 1, C]  
                                           mask=None,
                                           cross_attn_mask=None,
                                           seq_lens=seq_lens,                   # (CW) - for document masking in self-attention
                                           cross_seq_lens=seq_lens,             # (CW) - for document masking in cross-attention (with CR=1)
                                           tok_idx=tok_idx,                     # (CW) - pass in coarse time index for 1D RoPE
                                           cross_tok_idx=tok_idx,               # (CW) - pass in coarse time index for 1D RoPE (with CR=1)
                                           do_idx=do_idx, #.squeeze(0) if do_idx is not None else None,  # indices of dropped-out channels (CW)
        )

        return dec_out, enc_losses, dec_losses


    @torch.no_grad()
    def sample(self, encoder_input: torch.Tensor, seq_lens: torch.Tensor, tok_idx: torch.Tensor, sample_steps: int = 50, cfg: float = 1.0):

        device = encoder_input.device
        dtype = torch.bfloat16 # if device.type == "cuda" else torch.float16 # torch.float32
        # CPU Autocast only supports dtypes of torch.bfloat16, torch.float16 currently.
        with torch.autocast(device.type, dtype=dtype):

            do_idx = (encoder_input.sum(axis=2)==0).squeeze(0) # indices of dropped-out channels (CW) 

            # do_idx = None # [Set do_idx to None to disable printing of activation stats comparing channel drop-out]
            enc_out, _ = self.encoder(
                            token_values=encoder_input, 
                            seq_lens=seq_lens,
                            tok_idx=tok_idx,
                            do_idx=do_idx,
            )
            # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

            bsz, seqlen, dim = enc_out.shape
            dt_time = torch.tensor([1.0 / sample_steps] * bsz, device=enc_out.device).view(-1)


            z = self.global_sigma*torch.randn_like(encoder_input).to(enc_out.device) # init to rand
            # z = torch.zeros_like(encoder_input).to(enc_out.device) # init to zeros

            # Do not noise channel {x,y,z}-position in eeg_signal
            if self.dont_noise_chan_xyz:
                if dim==131 or dim==35:
                    z[:,:,:3] = encoder_input[:,:,:3]
                else:
                    pass
                    # print("NOTE: EEG channel {x,y,z}-position was never concatenated into signal.")
                    # import IPython; print('\n\nDebug:'); IPython.embed(); import time;  time.sleep(0.3)


            dt = dt_time.unsqueeze(-1).unsqueeze(-1)

            outputs = []
            for i in range(sample_steps, 0, -1):
                t = dt_time * i
                t_model = t.unsqueeze(1).unsqueeze(1)

                vc, _ = self.decoder(tokens=z.unsqueeze(1),
                                     cross_attended=enc_out, 
                                     timeD=t_model, 
                                     seq_lens=seq_lens,                   # for document masking in self-attention
                                     cross_seq_lens=seq_lens,             # for document masking in cross-attention (with CR=1)
                                     tok_idx=tok_idx,                     
                                     cross_tok_idx=tok_idx,                        
                )

                if cfg != 1.0:
                    vc_uncond, _ = self.decoder(tokens=z.unsqueeze(1),
                                                cross_attended=torch.zeros_like(enc_out), 
                                                timeD=t_model, 
                                                seq_lens=seq_lens,                          # for document masking in self-attention
                                                cross_seq_lens=seq_lens,                    # for document masking in cross-attention (with CR=1)
                                                tok_idx=tok_idx,                           
                                                cross_tok_idx=tok_idx,                                
                    )

                    vc = vc_uncond + cfg * (vc - vc_uncond) # starts at unconditioned, moves toward conditioned as cfg increases
                    
                z = z - dt * vc

                # Do not noise channel {x,y,z}-position in eeg_signal
                if self.dont_noise_chan_xyz:
                    if dim==131 or dim==35:
                        z[:,:,:3] = encoder_input[:,:,:3]
                    else:
                        # print("NOTE: EEG channel {x,y,z}-position was never concatenated into signal.")
                        # import IPython; print('\n\nDebug:'); IPython.embed(); import time;  time.sleep(0.3)
                        pass

                outputs.append(z)
            
            return z, outputs



    def reset_parameters(self):
        self.encoder.reset_parameters()
        self.decoder.reset_parameters()

    def init_weights(self):
        self.encoder.init_weights()
        self.decoder.init_weights()




# Optional policy for activation checkpointing. With None, we stick to the default (defined distributed.py: default_no_recompute_ops)
def get_no_recompute_ops():
    return None


# Optional and only used for fully shard options (fsdp) is choose. Highly recommanded for large models
def build_fsdp_grouping_plan(model_args: DecoderTransformerArgs) -> List[Tuple[str, bool]]:
    group_plan: List[Tuple[str, bool]] = []

    # # 1. Encoder Input
    group_plan.append(("encoder.output", False)) # <-- Changed to True
    group_plan.append(("decoder.output", False)) # Final output for main loss
    if model_args.decoder_repa_index != inf:
        group_plan.append(("decoder.repa_proj", False))
    if model_args.encoder_repa_index != inf:
        group_plan.append(("encoder.repa_proj", False))

    # 2. Encoder Transformer Blocks
    for i in range(model_args.n_layers):
        group_plan.append((f"encoder.layers.{i}", False))

    # 3. Decoder Transformer Blocks
    for i in range(model_args.n_layers):
        group_plan.append((f"decoder.layers.{i}", False))

    # 4. Add Decoder and Encoder themselves
    group_plan.append(("encoder", False))
    group_plan.append(("decoder", False))

    return group_plan