Self-contained HF-compatible ZUNA (vendored arch, byte-identical weights)

README.md

@@ -0,0 +1,52 @@
+---
+license: other
+tags:
+- eeg
+- eeg-foundation-model
+- neurotechx
+- zuna
+library_name: transformers
+pipeline_tag: feature-extraction
+---
+
+# ZUNA — EEG Foundation Model (NeuroTechX HF-compatible build)
+
+**ZUNA** is the [Zyphra](https://www.zyphra.com/) EEG foundation model — an
+encoder–decoder architecture (encoder → latent embeddings; flow-matching diffusion
+decoder), **382 M parameters**. This is a **self-contained, HuggingFace
+`trust_remote_code` build** prepared by **NeuroTechX** so the model loads anywhere with
+only `torch` + `transformers` + a few public pip deps — no private packages required.
+
+Part of the [NeuroTechX EEG-FM Collection](https://huggingface.co/collections/NeuroTechX/eeg-foundation-models-6a443a79e9a158d702f2e8e1)
+and the 2026 NeuroTechX DL/FM-on-EEG white paper (the successor to Roy et al. 2019).
+
+## Install
+```bash
+pip install torch transformers vector_quantize_pytorch einx einops
+```
+
+## Usage
+```python
+from transformers import AutoModel
+model = AutoModel.from_pretrained("NeuroTechX/zuna", trust_remote_code=True).cuda().eval()
+
+# encoder-only latents (the typical EEG-FM use):
+latents = model.encode(encoder_input, seq_lens, ...)   # [B, L', encoder_output_dim]
+```
+Or, in the NeuroTechX stack: `make_hf_encoder("NeuroTechX/zuna")` (emeg-fm).
+
+## What this build changes
+A faithful mirror of the weights in
+[`mhough/zuna-base`](https://huggingface.co/mhough/zuna-base), made **standalone**. The
+original `modeling_zuna.py` imported a private `zuna` package, so it could not load via
+`trust_remote_code` for anyone without that environment. Here the `EncoderDecoder`
+architecture is **vendored** from the Zyphra source + Meta's
+[`lingua`](https://github.com/facebookresearch/lingua) (MIT), imports rewritten to be
+self-contained, and the train-only activation-probe stubbed to a no-op at inference.
+**Weights are byte-identical** to the original.
+
+## Attribution & license
+- **Architecture & weights:** Zyphra (ZUNA).
+- **Vendored framework:** Meta `lingua` (MIT).
+- Licensing of the ZUNA weights should be confirmed with Zyphra; this build is offered
+  in good faith for open EEG-FM research and reproducibility.

bottlenecks.py

@@ -0,0 +1,18 @@
+from torch import Tensor
+import torch
+
+def imq_kernel2(X: Tensor, Y: Tensor, C: float) -> Tensor:
+    dist_sq = torch.cdist(X, Y, p=2).pow(2)
+    return C / (C + dist_sq)
+
+
+@torch.compile()
+def mmd_imq(X: Tensor, Y: Tensor, C: float) -> Tensor:
+    K_XX = imq_kernel2(X, X, C)
+    K_YY = imq_kernel2(Y, Y, C)
+    K_XY = imq_kernel2(X, Y, C)
+    n, m = X.size(0), Y.size(0)
+    term1 = (K_XX.sum().double() - K_XX.diag().sum().double()).double() / (n*(n-1)) if n>1 else 0.0
+    term2 = (K_YY.sum().double() - K_YY.diag().sum().double()).double() / (m*(m-1)) if m>1 else 0.0
+    term3 = 2 * K_XY.mean()
+    return (term1 + term2 - term3).float()

config.json

@@ -0,0 +1,60 @@
+{
+  "adaptive_loss_weighting": true,
+  "architectures": [
+    "ZunaModel"
+  ],
+  "bottleneck_type": "mmd",
+  "codebook_size": 1024,
+  "compression_free_conv_stem": false,
+  "decoder_encoder_dropout": 0.1,
+  "decoder_repa_index": {
+    "__float__": "Infinity"
+  },
+  "decoder_timestep_dropout": 0.1,
+  "dim": 1024,
+  "distill_output_dim": 0,
+  "dont_noise_chan_xyz": false,
+  "dropout_type": "zeros",
+  "dtype": "float32",
+  "encoder_hidden_dim": null,
+  "encoder_input_dim": 32,
+  "encoder_latent_downsample_factor": 1,
+  "encoder_output_dim": 32,
+  "encoder_repa_index": {
+    "__float__": "Infinity"
+  },
+  "encoder_sliding_window": 65536,
+  "ffn_dim_multiplier": null,
+  "head_dim": 64,
+  "huber_c": null,
+  "init_base_std": 0.02,
+  "init_std_factor": "disabled",
+  "input_dim": 32,
+  "learnable_bias": false,
+  "levels": [],
+  "max_seqlen": 50,
+  "model_type": "zuna",
+  "multiple_of": 256,
+  "n_heads": 8,
+  "n_kv_heads": null,
+  "n_layers": 16,
+  "norm_eps": 1e-05,
+  "num_fine_time_pts": 32,
+  "repa_dim": 1024,
+  "repa_loss_fn": "cosine",
+  "rope_dim": 4,
+  "rope_theta": 10000.0,
+  "seed": 42,
+  "seqlen_t": false,
+  "sliding_window": 65536,
+  "stft_global_sigma": 0.1,
+  "t_dim": 64,
+  "tok_idx_type": "{x,y,z,tc}",
+  "transformers_version": "5.2.0",
+  "weight_tying": false,
+  "xattn_sliding_window": 65536,
+  "auto_map": {
+    "AutoModel": "modeling_zuna.ZunaModel",
+    "AutoConfig": "configuration_zuna.ZunaConfig"
+  }
+}

configuration_zuna.py

@@ -0,0 +1,176 @@
+"""
+ZunaConfig: HuggingFace PretrainedConfig wrapper for the Zyphra ZUNA foundation model.
+
+Maps all fields from DecoderTransformerArgs (and its parent chain
+BaseTransformerArgs -> DecoderArgs -> DecoderTransformerArgs) to a standard
+HF config, so the model can be loaded with AutoConfig / trust_remote_code=True.
+
+Field sources:
+  BaseTransformerArgs  – lingua/transformer.py
+  DecoderArgs          – xattn.py
+  DecoderTransformerArgs – transformer.py (AY2latent_bci)
+"""
+
+from typing import List, Optional, Union
+from transformers import PretrainedConfig
+
+
+class ZunaConfig(PretrainedConfig):
+    model_type = "zuna"
+
+    def __init__(
+        self,
+        # ── BaseTransformerArgs ──────────────────────────────────────────────
+        dim: int = 1024,
+        n_layers: int = 10,
+        head_dim: Optional[int] = None,
+        n_heads: int = 8,
+        n_kv_heads: Optional[int] = None,
+        ffn_dim_multiplier: Optional[float] = None,
+        multiple_of: int = 256,
+        norm_eps: float = 1e-5,
+        rope_theta: float = 10000.0,
+        init_base_std: Optional[float] = 0.02,
+        init_std_factor: str = "disabled",
+        max_seqlen: int = 1024,
+        rope_dim: int = 1,          # 0=NoPE, 1=1D-RoPE, 4=4D-RoPE
+        tok_idx_type: Optional[str] = "t_coarse",
+        # ── DecoderArgs ─────────────────────────────────────────────────────
+        t_dim: int = 64,
+        seqlen_t: bool = False,
+        # ── DecoderTransformerArgs ───────────────────────────────────────────
+        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: Optional[int] = None,   # defaults to input_dim at runtime
+        encoder_output_dim: Optional[int] = None,  # defaults to input_dim*2 at runtime
+        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: Optional[List[int]] = None,
+        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,
+        **kwargs,
+    ):
+        # ── BaseTransformerArgs ──────────────────────────────────────────────
+        self.dim = dim
+        self.n_layers = n_layers
+        self.head_dim = head_dim
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.ffn_dim_multiplier = ffn_dim_multiplier
+        self.multiple_of = multiple_of
+        self.norm_eps = norm_eps
+        self.rope_theta = rope_theta
+        self.init_base_std = init_base_std
+        self.init_std_factor = init_std_factor
+        self.max_seqlen = max_seqlen
+        self.rope_dim = rope_dim
+        self.tok_idx_type = tok_idx_type
+        # ── DecoderArgs ─────────────────────────────────────────────────────
+        self.t_dim = t_dim
+        self.seqlen_t = seqlen_t
+        # ── DecoderTransformerArgs ───────────────────────────────────────────
+        self.seed = seed
+        self.weight_tying = weight_tying
+        self.sliding_window = sliding_window
+        self.xattn_sliding_window = xattn_sliding_window
+        self.input_dim = input_dim
+        self.decoder_encoder_dropout = decoder_encoder_dropout
+        self.decoder_timestep_dropout = decoder_timestep_dropout
+        self.encoder_sliding_window = encoder_sliding_window
+        self.encoder_input_dim = encoder_input_dim if encoder_input_dim is not None else input_dim
+        self.encoder_output_dim = encoder_output_dim if encoder_output_dim is not None else input_dim * 2
+        self.encoder_latent_downsample_factor = encoder_latent_downsample_factor
+        self.encoder_hidden_dim = encoder_hidden_dim
+        self.adaptive_loss_weighting = adaptive_loss_weighting
+        self.num_fine_time_pts = num_fine_time_pts
+        self.dont_noise_chan_xyz = dont_noise_chan_xyz
+        self.stft_global_sigma = stft_global_sigma
+        self.dropout_type = dropout_type
+        self.bottleneck_type = bottleneck_type
+        self.distill_output_dim = distill_output_dim
+        self.codebook_size = codebook_size
+        self.levels = levels if levels is not None else []
+        self.learnable_bias = learnable_bias
+        self.huber_c = huber_c
+        self.decoder_repa_index = decoder_repa_index
+        self.encoder_repa_index = encoder_repa_index
+        self.repa_dim = repa_dim
+        self.repa_loss_fn = repa_loss_fn
+        self.compression_free_conv_stem = compression_free_conv_stem
+
+        super().__init__(**kwargs)
+
+    def to_decoder_transformer_args(self):
+        """
+        Convert back to a DecoderTransformerArgs dataclass instance so the raw
+        Zyphra EncoderDecoder can be instantiated.
+        """
+        from .transformer import (
+            DecoderTransformerArgs,
+        )
+
+        return DecoderTransformerArgs(
+            dim=self.dim,
+            n_layers=self.n_layers,
+            head_dim=self.head_dim,
+            n_heads=self.n_heads,
+            n_kv_heads=self.n_kv_heads,
+            ffn_dim_multiplier=self.ffn_dim_multiplier,
+            multiple_of=self.multiple_of,
+            norm_eps=self.norm_eps,
+            rope_theta=self.rope_theta,
+            init_base_std=self.init_base_std,
+            init_std_factor=self.init_std_factor,
+            max_seqlen=self.max_seqlen,
+            rope_dim=self.rope_dim,
+            tok_idx_type=self.tok_idx_type,
+            t_dim=self.t_dim,
+            seqlen_t=self.seqlen_t,
+            seed=self.seed,
+            weight_tying=self.weight_tying,
+            sliding_window=self.sliding_window,
+            xattn_sliding_window=self.xattn_sliding_window,
+            input_dim=self.input_dim,
+            decoder_encoder_dropout=self.decoder_encoder_dropout,
+            decoder_timestep_dropout=self.decoder_timestep_dropout,
+            encoder_sliding_window=self.encoder_sliding_window,
+            encoder_input_dim=self.encoder_input_dim,
+            encoder_output_dim=self.encoder_output_dim,
+            encoder_latent_downsample_factor=self.encoder_latent_downsample_factor,
+            encoder_hidden_dim=self.encoder_hidden_dim,
+            adaptive_loss_weighting=self.adaptive_loss_weighting,
+            num_fine_time_pts=self.num_fine_time_pts,
+            dont_noise_chan_xyz=self.dont_noise_chan_xyz,
+            stft_global_sigma=self.stft_global_sigma,
+            dropout_type=self.dropout_type,
+            bottleneck_type=self.bottleneck_type,
+            distill_output_dim=self.distill_output_dim,
+            codebook_size=self.codebook_size,
+            levels=list(self.levels),
+            learnable_bias=self.learnable_bias,
+            huber_c=self.huber_c,
+            decoder_repa_index=self.decoder_repa_index,
+            encoder_repa_index=self.encoder_repa_index,
+            repa_dim=self.repa_dim,
+            repa_loss_fn=self.repa_loss_fn,
+            compression_free_conv_stem=self.compression_free_conv_stem,
+        )

conv_stem.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional
+import math
+
+class CausalConv2DStem(nn.Module):
+    """
+    A 'beefier' Causal 2D Convolutional Stem using standard convolutions.
+
+    Processes input [B, T, C=2*F] -> [B, T, F_out]. Uses standard Conv2D layers
+    for increased parameters. Maintains strict time causality via manual padding.
+    Uses PyTorch's truncated normal initialization.
+    Structure: Pointwise -> Act -> Conv2D -> Act -> Conv2D -> Act -> Pointwise
+
+    Args:
+        input_features (int): Input features C (even, = 2*F).
+        hidden_channels (int): Intermediate convolution channels.
+        time_kernel_size (int): Kernel size for time dim (>= 1).
+        freq_kernel_size (int): Kernel size for freq dim (>= 1, typically odd).
+        compress_channels (bool): If True, output F_out = F. If False, F_out = 2*F.
+        activation (nn.Module): Activation function. Default: nn.GELU().
+        init_std (float): Std dev for truncated normal weight init. Default: 0.02.
+    """
+    def __init__(
+        self,
+        input_features: int,
+        hidden_channels: int,
+        time_kernel_size: int = 3,
+        freq_kernel_size: int = 3,
+        compress_channels: bool = True,
+        activation: Optional[nn.Module] = None,
+    ):
+        super().__init__()
+        if not (isinstance(input_features, int) and input_features > 0 and input_features % 2 == 0):
+            raise ValueError("input_features must be an even positive integer.")
+        if not (isinstance(hidden_channels, int) and hidden_channels > 0):
+            raise ValueError("hidden_channels must be a positive integer.")
+        if not (isinstance(time_kernel_size, int) and time_kernel_size >= 1):
+            raise ValueError("time_kernel_size must be >= 1.")
+        if not (isinstance(freq_kernel_size, int) and freq_kernel_size >= 1):
+             raise ValueError("freq_kernel_size must be >= 1.")
+
+        self.input_features = input_features
+        self.freq_dim = input_features // 2
+        self.target_channels = 1 if compress_channels else 2
+        self._activation = activation() if activation is not None else nn.GELU()
+
+        self._time_pad_left = time_kernel_size - 1
+        self._freq_pad_sym = (freq_kernel_size - 1) // 2
+        self._causal_padding = (self._freq_pad_sym, self._freq_pad_sym, self._time_pad_left, 0)
+
+        self._pointwise1 = nn.Conv2d(2, hidden_channels, 1, 1, bias=False)
+        self._conv1 = nn.Conv2d(hidden_channels, hidden_channels, (time_kernel_size, freq_kernel_size), 1, padding=0, bias=False)
+        self._conv2 = nn.Conv2d(hidden_channels, hidden_channels, (time_kernel_size, freq_kernel_size), 1, padding=0, bias=False)
+        self._pointwise2 = nn.Conv2d(hidden_channels, self.target_channels, 1, 1, bias=True)
+        self.output_features = self.target_channels * self.freq_dim
+
+
+    def reset_parameters(self, std: float):
+        """Initialize Conv2d weights with truncated normal and biases with zeros."""
+        def init_the_shit(module):
+            if isinstance(module, nn.Conv2d):
+                nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-3*std, b=3*std)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+        self.apply(lambda module: init_the_shit(module))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """ Args: x [B, T, C=2*F]. Returns: [B, T, F_out]. """
+        B, T, C = x.shape
+        if C != self.input_features: raise ValueError(f"Input C={C} != expected {self.input_features}")
+
+        x = rearrange(x, 'b t (c f) -> b c t f', c=2, f=self.freq_dim)
+        x = self._activation(self._pointwise1(x))
+        x = F.pad(x, self._causal_padding)
+        x = self._activation(self._conv1(x))
+        x = F.pad(x, self._causal_padding)
+        x = self._activation(self._conv2(x))
+        x = self._pointwise2(x) # [B, target_channels, T, F]
+        x = rearrange(x, 'b c t f -> b t (c f)') # [B, T, F_out]
+        return x
+
+    def get_output_dim(self) -> int:
+        return self.output_features

lingua_transformer.py

@@ -0,0 +1,557 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union, Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn.attention.flex_attention import (
+    BlockMask,
+    flex_attention,
+    _mask_mod_signature,
+)
+
+class _NoProbe:  # inference no-op (Meta lingua activation-probe is train-only)
+    @staticmethod
+    def log_stats(x, name=None):
+        return x
+probe = _NoProbe()
+flex_attention_comp = torch.compile(flex_attention, dynamic=True, mode='max-autotune')
+
+class InitStdFactor(Enum):
+    DISABLED = "disabled"  # Init std is divided by 1.0
+    GLOBAL_DEPTH = "global_depth"  # Init std is divided by sqrt(2*n_layers)
+    CURRENT_DEPTH = "current_depth"  # Init std is divided by sqrt(2*depth)
+    DIM_RATIO = "dim_ratio"  # Init std is divided by model_dim/4096
+
+
+@dataclass
+class BaseTransformerArgs:
+
+    dim: int = 1024
+    n_layers: int = 10
+    head_dim: Optional[int] = None
+    n_heads: Optional[int] = None
+    n_kv_heads: Optional[int] = None
+
+    ffn_dim_multiplier: Optional[float] = None
+
+    multiple_of: int = 256
+
+    norm_eps: float = 1e-5
+
+    rope_theta: float = 10000.0
+
+    init_base_std: Optional[float] = None
+    init_std_factor: str = "disabled"
+
+    max_seqlen: int = 1024
+    rope_dim: int = 1 # 0 = NoPE, 1 = 1D-RoPE, 4 = 4D=RoPE.
+    tok_idx_type: Optional[str] = "t_coarse"
+
+def cross_entropy(pred, target, **kwargs):
+    return F.nll_loss(
+        F.log_softmax(pred.flatten(end_dim=-2).float(), -1),
+        target.flatten(end_dim=-1),
+        **kwargs,
+    )
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int, dim: int) -> torch.Tensor:
+    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
+    assert dim == 2, "Only dim=2 is supported. Check the implementation for other dims."
+    bs, slen, n_kv_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        x[:, :, :, None, :]
+        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
+        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
+    )
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    """
+    Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+    This function calculates a frequency tensor with complex exponentials using the given dimension 'dim'
+    and the end index 'end'. The 'theta' parameter scales the frequencies.
+    The returned tensor contains complex values in complex64 data type.
+
+    Args:
+        dim (int): Dimension of the frequency tensor.
+        end (int): End index for precomputing frequencies.
+        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
+
+    Returns:
+        torch.Tensor: Precomputed frequency tensor with complex exponentials.
+    """
+
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+
+    cos, sin = freqs.cos(), freqs.sin()
+
+    return torch.stack((cos, -sin, sin, cos), dim=-1).view(*freqs.size(), 2, 2)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor, seq_dim: int):
+    """
+    Reshape frequency tensor for broadcasting it with another tensor.
+
+    This function reshapes the frequency tensor to have the same shape as the target tensor 'x'
+    for the purpose of broadcasting the frequency tensor during element-wise operations.
+
+    Args:
+        freqs_cis (torch.Tensor): Frequency tensor to be reshaped.
+        x (torch.Tensor): Target tensor for broadcasting compatibility.
+        seq_dim (int): Sequence dimension index.
+
+    Returns:
+        torch.Tensor: Reshaped frequency tensor.
+    """
+
+    ndim = x.ndim
+    assert 0 <= seq_dim < ndim
+    assert freqs_cis.shape == (
+        x.shape[seq_dim],
+        x.shape[-3],
+        2,
+        2,
+    ), f"freqs_cis vs x: {(freqs_cis.shape, x.shape)}. freqs_cis should be{(x.shape[seq_dim], x.shape[-3], 2, 2)}."
+    shape = [
+        d if i == seq_dim or i == ndim - 3 else 1 for i, d in enumerate(x.shape[:-2])
+    ] + [2, 2]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    seq_dim: int,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    xq_ = xq.reshape(*xq.shape[:-1], -1, 1, 2)  # B S D -> B S D/2 1 2
+    xk_ = xk.reshape(*xk.shape[:-1], -1, 1, 2)  # B S D -> B S D/2 1 2
+
+    freqs_cis = reshape_for_broadcast(
+        freqs_cis, xq_, seq_dim
+    ).float()  # S D/2 2 2 -> 1 S 1 D/2 2 2
+    xq_out = (xq_ * freqs_cis).sum(5).flatten(3)
+    xk_out = (xk_ * freqs_cis).sum(5).flatten(3)
+
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+
+
+
+# Rotary embedding as in xformer
+class RotaryEmbedding(torch.nn.Module):
+    """
+    RotaryEmbedding Module
+    """
+
+    def __init__(self, theta: float, head_dim: int, max_seqlen: int = 1024, rope_dim: int = 1):
+        super().__init__()
+
+        self.theta = theta
+        self.head_dim = head_dim
+        self.max_seqlen = max_seqlen
+        self.rope_dim = rope_dim
+
+        assert head_dim % rope_dim == 0, f"head_dim must be divisible by rope_dim, got {head_dim} and {rope_dim}"
+
+        self.register_buffer(
+            "freqs_cis",
+            precompute_freqs_cis(dim=head_dim//rope_dim, end=max_seqlen, theta=theta),
+            persistent=False,
+        )
+
+    def reset_parameters(self):
+        self.freqs_cis[...] = precompute_freqs_cis(
+            dim=self.head_dim//self.rope_dim, end=self.max_seqlen, theta=self.theta
+        )
+
+    def forward(
+        self, seqlen: Optional[int] = None, tok_idx: Optional[torch.Tensor] = None
+    ):
+        """
+        Return freqs_cis corresponding to consecutive seqlen positions or the corresponding tok_idx positions
+        Args:
+            seqlen (int): Contiguous sequence length
+            tok_idx (torch.Tensor[int]): Position indices of each token. This overrides seqlen.
+
+        Returns:
+            Tuple(torch.Tensor, torch.Tensor): Embedded input tensor and freqs_cis
+        """
+
+        tok_idx = None # HARDCODE (CW)! SEE NOTE BELOW. WILL USE SEQLEN PATH.   
+
+        test = (seqlen is not None) or (tok_idx is not None)
+        assert test, "Should provide atleast seqlen or tok_idx"
+        if tok_idx is not None:
+            return self.freqs_cis[tok_idx] # NOTE: DONT WANT TO INDEX WITH TOK_IDX HERE AND THEN AGAIN INSIDE ATTENTION.FORWARD - DOUBLE DOING
+        elif seqlen is not None:
+            return self.freqs_cis[0:seqlen]
+
+
+class RMSNorm(nn.Module):
+    """
+    Initialize the RMSNorm normalization layer.
+
+    Args:
+        dim (int): The dimension of the input tensor.
+        eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+    Attributes:
+        eps (float): A small value added to the denominator for numerical stability.
+        weight (nn.Parameter): Learnable scaling parameter.
+
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6, channel_dim=-1):
+        super().__init__()
+        self.eps = eps
+        self.channel_dim = channel_dim
+        if channel_dim != -1: #channel_dim is the index of the channel dimension, dim is the number of channels. assume 4 dimensions.
+            self.weight = nn.Parameter(torch.ones([1]*channel_dim + [dim] + [1]*(4-channel_dim-1)))
+        else:
+            self.weight = nn.Parameter(torch.ones(dim))
+
+
+    def _norm(self, x: torch.Tensor):
+        return x * torch.rsqrt((x * x).mean(self.channel_dim, keepdim=True) + self.eps)
+
+    def forward(self, x: torch.Tensor, y: Optional[torch.Tensor] = None) -> torch.Tensor:
+        x = probe.log_stats(x, "resid")
+        output = self._norm(x.float())
+        return (output * self.weight.float()).type_as(x)
+
+    def reset_parameters(self):
+        torch.nn.init.ones_(self.weight)  # type: ignore
+
+class TiedLinear(nn.Module):
+    def __init__(self, tied_module: nn.Module) -> None:
+        super().__init__()
+        self.tied_module = tied_module
+        if not hasattr(tied_module, "weight"):
+            raise AttributeError(
+                "Provided module does not have attribute 'weight'. Please check your tied_module."
+            )
+
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
+        return F.linear(x, self.tied_module.weight)
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        head_dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        rope_theta: float,
+        rope_dim: int,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.head_dim = head_dim
+        self.rope_theta = rope_theta
+        self.rope_dim = rope_dim
+
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.heads_per_group = self.n_heads // self.n_kv_heads
+
+        self.wq = nn.Linear(
+            dim,
+            n_heads * head_dim,
+            bias=False,
+        )
+        self.wk = nn.Linear(
+            dim,
+            n_kv_heads * head_dim,
+            bias=False,
+        )
+        self.wv = nn.Linear(
+            dim,
+            n_kv_heads * head_dim,
+            bias=False,
+        )
+
+        self.wo = nn.Linear(
+            n_heads * head_dim,
+            dim,
+            bias=False,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freq_cis: torch.Tensor,
+        tok_idx: Optional[torch.Tensor] = None,
+        mask: Optional[Union[BlockMask,  str]] = None,
+        attn_impl: str = "sdpa",
+    ) -> torch.Tensor:
+
+        # B S D
+        bsz, seq_len, dim = x.shape
+        xq = self.wq(x.view_as(x)) 
+        xk = self.wk(x.view_as(x))
+        xv = self.wv(x.view_as(x))
+
+        output_shape = xq.shape
+        # B S D -> B S H Dh         (where D = H*Dh)
+        xq = xq.view(bsz, seq_len, self.n_heads, self.head_dim)
+        xk = xk.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
+        xv = xv.view(bsz, seq_len, self.n_kv_heads, self.head_dim)
+
+        if self.rope_dim==0:
+            pass
+        elif self.rope_dim==1:
+            if tok_idx is not None:
+                xq, xk = apply_rotary_emb(xq, xk, 1, freq_cis[tok_idx])     # this edit mirrors what is inside RotaryEmbedding class. To use tok_idx
+            else:
+                xq, xk = apply_rotary_emb(xq, xk, 1, freq_cis[0:seq_len])   # This is how it was. (SEEMS TO ASSUME WE ARE USING MAX_SEQLEN, NOT TOK_IDX)
+        elif self.rope_dim==4:
+            freqcis_parts = []
+            for i in range(self.rope_dim):
+                freqcis_parts.append(freq_cis[tok_idx[:, i]])
+            freqcis_4RoPE = torch.cat(freqcis_parts, dim=1)
+
+            # Now apply 4D-axial-RoPE
+            xq, xk = apply_rotary_emb(xq, xk, 1, freqcis_4RoPE) 
+
+
+        else:
+            print(f"I dont know how to handle {self.rope_dim=} inside lingua.transformer.Attention.forward")
+            import IPython; print('\n\nDebug:'); IPython.embed(); import time;  time.sleep(0.3)
+
+
+        # print(x, xq, xk, xv, freq_cis, tok_idx, seq_len)
+
+        # This condition helps us be easily compatible
+        # with inference by adding a pluggable KVCache
+        if hasattr(self, "kv_cache"):
+            xk, xv = self.kv_cache.update(xk, xv, tok_idx)
+
+        xk = repeat_kv(xk, self.heads_per_group, dim=2)
+        xv = repeat_kv(xv, self.heads_per_group, dim=2)
+
+        if attn_impl == "flex_attention":
+            assert mask is None or isinstance(mask, BlockMask)
+            xq, xk, xv = map(lambda e: e.transpose(1, 2), (xq, xk, xv))
+            if xq.device.type == "mps":
+                # MPS does not support flex_attention; fall back to SDPA with dense mask
+                if mask is not None:
+                    S = xq.shape[2]
+                    q_idx = torch.arange(S, device='cpu')
+                    kv_idx = torch.arange(S, device='cpu')
+                    dense_bool = mask.mask_mod(0, 0, q_idx.unsqueeze(1), kv_idx.unsqueeze(0))
+                    attn_mask = torch.zeros(1, 1, S, S, dtype=xq.dtype, device=xq.device)
+                    attn_mask.masked_fill_(~dense_bool.unsqueeze(0).unsqueeze(0).to(xq.device), float("-inf"))
+                else:
+                    attn_mask = None
+                output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask)
+            elif xq.device.type == "cuda":
+                output = flex_attention_comp(xq, xk, xv, block_mask=mask)
+            else:
+                output = flex_attention(xq, xk, xv, block_mask=mask)
+            output = output.transpose(1, 2).contiguous()  # B H S D -> B S H D
+
+        elif attn_impl == "sdpa":
+            xq, xk, xv = map(lambda e: e.transpose(1, 2), (xq, xk, xv))
+            assert mask is None or isinstance(mask, (str, torch.Tensor))
+            is_causal = (mask == "causal") if isinstance(mask, str) else False
+            mask = mask if isinstance(mask, torch.Tensor) else None
+            output = F.scaled_dot_product_attention(
+                xq,
+                xk,
+                xv,
+                is_causal=is_causal,
+                attn_mask=mask,
+            )
+            output = output.transpose(1, 2).contiguous()  # B H S D -> B S H D
+        else:
+            raise NotImplementedError(
+                f"Attention implementation {attn_impl} not supported"
+            )
+
+        output = self.wo(output.reshape(output_shape))
+
+        return output
+
+    def reset_parameters(self, init_std=None, factor=1.0):
+        init_std = init_std or (self.dim ** (-0.5))
+
+        for w in [self.wq, self.wk, self.wv]:
+            nn.init.trunc_normal_(
+                w.weight,
+                mean=0.0,
+                std=init_std,
+                a=-3 * init_std,
+                b=3 * init_std,
+            )
+
+        nn.init.trunc_normal_(
+            self.wo.weight,
+            mean=0.0,
+            std=init_std / factor,
+            a=-3 * init_std,
+            b=3 * init_std,
+        )
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        multiple_of: int,
+        ffn_dim_multiplier: Optional[float],
+        mp_size: int = 1,
+    ):
+        super().__init__()
+
+        hidden_dim = int(2 * hidden_dim / 3)
+        if ffn_dim_multiplier is not None:
+            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        assert hidden_dim % mp_size == 0
+
+        self.dim = dim
+        self.hidden_dim = hidden_dim
+
+        self.w1 = nn.Linear(
+            dim,
+            hidden_dim,
+            bias=False,
+        )
+        self.w3 = nn.Linear(
+            dim,
+            hidden_dim,
+            bias=False,
+        )
+        self.w2 = nn.Linear(
+            hidden_dim,
+            dim,
+            bias=False,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # B S D
+        x1 = self.w1(x.view_as(x))
+        x3 = self.w3(x.view_as(x))
+        output = self.w2(F.silu(x1) * x3)
+        return output
+
+    def reset_parameters(self, init_std=None, factor=1.0):
+        in_init_std = init_std or (self.dim ** (-0.5))
+        out_init_std = init_std or (self.hidden_dim ** (-0.5))
+        in_init_std = in_init_std
+        out_init_std = out_init_std / factor
+        for w in [self.w1, self.w3]:
+            nn.init.trunc_normal_(
+                w.weight,
+                mean=0.0,
+                std=in_init_std,
+                a=-3 * in_init_std,
+                b=3 * in_init_std,
+            )
+        nn.init.trunc_normal_(
+            self.w2.weight,
+            mean=0.0,
+            std=out_init_std,
+            a=-3 * out_init_std,
+            b=3 * out_init_std,
+        )
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: BaseTransformerArgs):
+        super().__init__()
+
+        assert (args.head_dim is not None) or (
+            args.n_heads is not None
+        ), "Should specify at least head_dim or n_heads"
+        self.head_dim = args.head_dim or args.dim // args.n_heads
+        self.n_heads = args.n_heads or args.dim // args.head_dim
+        self.n_kv_heads = args.n_kv_heads or self.n_heads
+
+        assert args.n_heads % self.n_kv_heads == 0
+        assert args.dim % args.n_heads == 0
+
+        self.attention = Attention(
+            dim=args.dim,
+            head_dim=self.head_dim,
+            n_heads=self.n_heads,
+            n_kv_heads=self.n_kv_heads,
+            rope_theta=args.rope_theta,
+            rope_dim=args.rope_dim,
+        )
+        self.feed_forward = FeedForward(
+            dim=args.dim,
+            hidden_dim=4 * args.dim,
+            multiple_of=args.multiple_of,
+            ffn_dim_multiplier=args.ffn_dim_multiplier,
+        )
+        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
+        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        freq_cis: torch.Tensor,
+        tok_idx: Optional[torch.Tensor] = None,
+        mask: Optional[Union[BlockMask,  str]] = None,
+        attn_impl: str = "sdpa",
+        do_idx: Optional[torch.Tensor] = None,
+        print_layerwise_activation_stats: bool = False,
+    ) -> torch.Tensor:
+
+
+        if print_layerwise_activation_stats and do_idx is not None:
+
+            # Print all the activation stats for the dropped and non-dropped tokens if do_idx is provided
+            x_normed = self.attention_norm(x)
+            print(f"\n\t Encoder attn_norm (drop-out): mean={x[:, do_idx, :].mean().item():.6f}, std={x[:, do_idx, :].std().item():.6f}", end=" --> ")
+            print(f"mean={x_normed[:, do_idx, :].mean().item():.6f}, std={x_normed[:, do_idx, :].std().item():.6f}") 
+            print(f"\t Encoder attn_norm (non-drop): mean={x[:, ~do_idx, :].mean().item():.6f}, std={x[:, ~do_idx, :].std().item():.6f}", end=" --> ")
+            print(f"mean={x_normed[:, ~do_idx, :].mean().item():.6f}, std={x_normed[:, ~do_idx, :].std().item():.6f}") 
+            h = x + self.attention(                         #  lingua.transformer.Attention
+                x_normed,
+                freq_cis,
+                tok_idx=tok_idx,
+                mask=mask, 
+                attn_impl=attn_impl,
+            )
+            h_normed = self.ffn_norm(h) 
+            print(f"\n\t Encoder ffn_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"\t Encoder ffn_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}") 
+            out = h + self.feed_forward(h_normed)  #  lingua.transformer.FeedForward
+
+        else:
+            h = x + self.attention(                         
+                self.attention_norm(x),
+                freq_cis,
+                tok_idx=tok_idx,
+                mask=mask, 
+                attn_impl=attn_impl,
+            )
+            out = h + self.feed_forward(self.ffn_norm(h))  
+
+        return out
+
+    def init_weights(self, init_std=None, factor=1.0):
+        self.attention.reset_parameters(init_std, factor)
+        self.attention_norm.reset_parameters()
+
+        self.feed_forward.reset_parameters(init_std, factor)
+        self.ffn_norm.reset_parameters()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16f00d82e6ac5d42cfdb9d153850d114e8b280fb9dec79ea1de1b9eac03d1271
+size 1528130680

modeling_zuna.py

@@ -0,0 +1,211 @@
+"""
+ZunaModel: HuggingFace PreTrainedModel wrapper for the Zyphra ZUNA foundation model.
+
+Architecture (from Zyphra source):
+  EncoderDecoder
+  ├── encoder: EncoderTransformer   → produces latent embeddings [B, L', encoder_output_dim]
+  └── decoder: DecoderTransformer   → flow-matching diffusion decoder
+
+The primary user-facing method is encode(), which runs only the encoder and
+returns the latent embeddings.  forward() mirrors the Zyphra sample() signature
+for full encode→decode inference.
+
+Encoding call (verified against eeg_eval.py and EncoderDecoder.sample()):
+
+    do_idx = (encoder_input.sum(axis=2) == 0).squeeze(0)
+    enc_out, _ = model.model.encoder(
+        token_values=encoder_input,
+        seq_lens=seq_lens,
+        tok_idx=tok_idx,
+        do_idx=do_idx,
+    )
+    # enc_out: [B, L/downsample_factor, encoder_output_dim]
+
+State-dict convention: the safetensors file ships keys prefixed with "model."
+which must be stripped before load_state_dict (see convert_weights.py).
+After stripping, keys match EncoderDecoder's submodule names directly:
+  encoder.tok_embeddings.weight, decoder.layers.0.*, etc.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from .configuration_zuna import ZunaConfig
+
+
+class ZunaModel(PreTrainedModel):
+    config_class = ZunaConfig
+    # The raw Zyphra weights match EncoderDecoder's attribute tree directly
+    # (after "model." prefix is stripped in convert_weights.py).
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config: ZunaConfig):
+        super().__init__(config)
+
+        from .transformer import (
+            EncoderDecoder,
+        )
+
+        args = config.to_decoder_transformer_args()
+        self.model = EncoderDecoder(args)
+
+        # Store config fields needed at inference time.
+        self.tok_idx_type = config.tok_idx_type
+        self.rope_dim = config.rope_dim
+
+        self.post_init()
+
+    # ── internal helper ────────────────────────────────────────────────────
+
+    def _build_tok_idx(
+        self,
+        encoder_input: torch.Tensor,
+        seq_lens: torch.Tensor,
+        t_coarse: Optional[torch.Tensor],
+        chan_id: Optional[torch.Tensor],
+        chan_pos_discrete: Optional[torch.Tensor],
+    ) -> Optional[torch.Tensor]:
+        """
+        Replicate the tok_idx construction from EncoderDecoder.forward()
+        (transformer.py lines 1001-1012).  Only the types used by the public
+        ZUNA checkpoint are implemented here; extend as needed.
+        """
+        if self.tok_idx_type is None:
+            return None
+        elif self.tok_idx_type == "t_coarse" and self.rope_dim == 1:
+            if t_coarse is None:
+                raise ValueError("tok_idx_type='t_coarse' requires t_coarse tensor")
+            return t_coarse
+        elif self.tok_idx_type == "chan_id" and self.rope_dim == 1:
+            if chan_id is None:
+                raise ValueError("tok_idx_type='chan_id' requires chan_id tensor")
+            return chan_id
+        elif self.tok_idx_type == "stack_arange_seqlen" and self.rope_dim == 1:
+            return torch.hstack(
+                [torch.arange(sl) for sl in seq_lens]
+            ).unsqueeze(0).unsqueeze(-1)
+        elif self.tok_idx_type == "{x,y,z,tc}" and self.rope_dim == 4:
+            if chan_pos_discrete is None or t_coarse is None:
+                raise ValueError(
+                    "tok_idx_type='{x,y,z,tc}' requires chan_pos_discrete and t_coarse"
+                )
+            return torch.cat((chan_pos_discrete, t_coarse), dim=2)
+        else:
+            raise ValueError(
+                f"Unsupported tok_idx_type={self.tok_idx_type!r} with rope_dim={self.rope_dim}"
+            )
+
+    # ── encode ────────────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def encode(
+        self,
+        encoder_input: torch.Tensor,
+        seq_lens: torch.Tensor,
+        t_coarse: Optional[torch.Tensor] = None,
+        chan_id: Optional[torch.Tensor] = None,
+        chan_pos_discrete: Optional[torch.Tensor] = None,
+        tok_idx: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Run only the encoder and return latent embeddings.
+
+        Args:
+            encoder_input:    EEG signal tensor  [B, seqlen, input_dim]
+            seq_lens:         Sequence lengths   [B]  (used for document masking)
+            t_coarse:         Coarse time index  [B, seqlen, 1]  (needed when tok_idx_type='t_coarse')
+            chan_id:          Channel IDs        [B, seqlen, 1]  (needed when tok_idx_type='chan_id')
+            chan_pos_discrete: Discrete 3D positions [B, seqlen, 3] (needed for 4D-RoPE)
+            tok_idx:          Pre-built token index tensor; if provided, skips the
+                              automatic _build_tok_idx() construction above.
+
+        Returns:
+            Latent embeddings  [B, seqlen // downsample_factor, encoder_output_dim]
+        """
+        if encoder_input.ndim == 2:
+            encoder_input = encoder_input.unsqueeze(0)
+
+        if tok_idx is None:
+            tok_idx = self._build_tok_idx(
+                encoder_input, seq_lens, t_coarse, chan_id, chan_pos_discrete
+            )
+
+        # Identify dropped-out channels: columns that are all-zero
+        # (mirrors EncoderDecoder.sample() line 1055)
+        do_idx = (encoder_input.sum(axis=2) == 0).squeeze(0)
+
+        enc_out, _ = self.model.encoder(
+            token_values=encoder_input,
+            seq_lens=seq_lens,
+            tok_idx=tok_idx,
+            do_idx=do_idx,
+        )
+        return enc_out  # [B, L', encoder_output_dim]
+
+    # ── forward ───────────────────────────────────────────────────────────
+
+    def forward(
+        self,
+        encoder_input: torch.Tensor,
+        seq_lens: torch.Tensor,
+        t_coarse: Optional[torch.Tensor] = None,
+        chan_id: Optional[torch.Tensor] = None,
+        chan_pos_discrete: Optional[torch.Tensor] = None,
+        tok_idx: Optional[torch.Tensor] = None,
+        sample_steps: int = 50,
+        cfg: float = 1.0,
+        return_latents: bool = False,
+    ) -> BaseModelOutput:
+        """
+        Full encode → diffusion-decode pass, mirroring EncoderDecoder.sample().
+
+        Args:
+            encoder_input:    EEG signal tensor  [B, seqlen, input_dim]
+            seq_lens:         Sequence lengths   [B]
+            t_coarse:         Coarse time index  [B, seqlen, 1]
+            chan_id:          Channel IDs        [B, seqlen, 1]
+            chan_pos_discrete: Discrete 3D positions [B, seqlen, 3]
+            tok_idx:          Pre-built token index (overrides auto construction)
+            sample_steps:     Number of flow-matching diffusion steps (default 50)
+            cfg:              Classifier-free guidance scale (1.0 = no CFG)
+            return_latents:   If True, also return encoder latents in hidden_states
+
+        Returns:
+            BaseModelOutput with:
+                last_hidden_state: reconstructed signal  [B, seqlen, input_dim]
+                hidden_states:     (enc_out,) if return_latents else None
+        """
+        if encoder_input.ndim == 2:
+            encoder_input = encoder_input.unsqueeze(0)
+
+        if tok_idx is None:
+            tok_idx = self._build_tok_idx(
+                encoder_input, seq_lens, t_coarse, chan_id, chan_pos_discrete
+            )
+
+        reconstruction, _ = self.model.sample(
+            encoder_input=encoder_input,
+            seq_lens=seq_lens,
+            tok_idx=tok_idx,
+            sample_steps=sample_steps,
+            cfg=cfg,
+        )
+
+        hidden_states = None
+        if return_latents:
+            latents = self.encode(
+                encoder_input=encoder_input,
+                seq_lens=seq_lens,
+                tok_idx=tok_idx,
+            )
+            hidden_states = (latents,)
+
+        return BaseModelOutput(
+            last_hidden_state=reconstruction,
+            hidden_states=hidden_states,
+        )

transformer.py

@@ -0,0 +1,1169 @@
+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

utils.py

@@ -0,0 +1,169 @@
+import contextlib
+import math
+import os
+from collections.abc import Generator, Iterable
+from datetime import timedelta
+
+import torch
+import torch.distributed._functional_collectives as funcol
+import torch.distributed.distributed_c10d as c10d
+from torch import distributed as dist
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+
+import matplotlib.pyplot as plt
+from datetime import datetime
+import torch
+
+
+@torch.no_grad()
+def clip_grad_norm_(
+    parameters: torch.Tensor | Iterable[torch.Tensor],
+    max_norm: float,
+    norm_type: float = 2.0,
+    error_if_nonfinite: bool = False,
+    foreach: bool | None = None,
+    pp_mesh: DeviceMesh | None = None,
+) -> torch.Tensor:
+    """
+    Clip the gradient norm of an iterable of parameters.
+
+    Gradient norm clipping requires computing the gradient norm over the entire model.
+    `torch.nn.utils.clip_grad_norm_` only computes gradient norm along DP/FSDP/TP dimensions.
+    We need to manually reduce the gradient norm across PP stages.
+    See https://github.com/pytorch/torchtitan/issues/596 for details.
+
+    Args:
+        parameters: an iterable of Tensors or a single Tensor that will have gradients normalized
+        max_norm (float): max norm of the gradients
+        norm_type (float): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+        error_if_nonfinite (bool): if True, an error is thrown if the total
+            norm of the gradients from :attr:`parameters` is ``nan``,
+            ``inf``, or ``-inf``. Default: False (will switch to True in the future)
+        foreach (bool): use the faster foreach-based implementation.
+            If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently
+            fall back to the slow implementation for other device types.
+            Default: ``None``
+        pp_mesh: pipeline parallel device mesh. If not None, will reduce gradient norm across PP stages.
+
+    Returns:
+        Total norm of the parameter gradients (viewed as a single vector).
+
+    """
+    grads = [p.grad for p in parameters if p.grad is not None]
+    total_norm = torch.nn.utils.get_total_norm(
+        grads, norm_type, error_if_nonfinite, foreach
+    )
+
+    # If total_norm is a DTensor, the placements must be `torch.distributed._tensor.ops.math_ops._NormPartial`.
+    # We can simply reduce the DTensor to get the total norm in this tensor's process group
+    # and then convert it to a local tensor.
+    # NOTE: It has two purposes:
+    #       1. to make sure the total norm is computed correctly when PP is used (see below)
+    #       2. to return a reduced total_norm tensor whose .item() would return the correct value
+    if isinstance(total_norm, DTensor):
+        # Will reach here if any non-PP parallelism is used.
+        # If only using PP, total_norm will be a local tensor.
+
+        # Remove FT replicate dimension if it exists.
+        total_norm = total_norm.full_tensor()
+
+    total_norm **= norm_type
+    dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=pp_mesh.get_group())
+    total_norm **= 1.0 / norm_type
+
+    torch.nn.utils.clip_grads_with_norm_(parameters, max_norm, total_norm, foreach)
+    return total_norm
+
+
+def reconstruct_full_mask(mask):
+        """
+        Utilities for testing and visualizing attention masks in EEG transformer models.
+
+        Functions:
+        - reconstruct_full_mask: Converts sparse block attention mask to full dense mask
+        - visualize_attention_mask: Creates and saves attention pattern visualization
+        """
+
+        # Get the block structure (260x260 with 1s for active blocks)
+        block_structure = mask.to_dense()[0, 0]  # 260x260 sparse block pattern
+        full_shape = mask.shape[-1]  # Get the full sequence length (33180)
+        full_mask = torch.zeros(full_shape, full_shape, device=block_structure.device, dtype=torch.bool)
+
+        # Reconstruct full mask from active blocks only
+        block_size = mask.BLOCK_SIZE[0]  # 128
+        active_blocks = torch.where(block_structure == 1)  # Get coordinates of active blocks
+
+        for q_block, kv_block in zip(active_blocks[0], active_blocks[1]):
+            # Fill in the 128x128 regions for active blocks
+            q_start, q_end = q_block * block_size, min((q_block + 1) * block_size, full_shape)
+            kv_start, kv_end = kv_block * block_size, min((kv_block + 1) * block_size, full_shape)
+
+            # Use mask_mod to get actual within-block attention pattern (vectorized)
+            q_indices = torch.arange(q_start, q_end, device=block_structure.device)
+            kv_indices = torch.arange(kv_start, kv_end, device=block_structure.device)
+            q_grid, kv_grid = torch.meshgrid(q_indices, kv_indices, indexing='ij')
+
+            # Call mask_mod with flattened arrays for efficiency
+            block_mask = mask.mask_mod(0, 0, q_grid.flatten(), kv_grid.flatten())
+            block_mask = block_mask.reshape(q_end - q_start, kv_end - kv_start)
+
+            full_mask[q_start:q_end, kv_start:kv_end] = block_mask
+
+        return full_mask
+
+def visualize_attention_mask(mask, sample_size=5000, title_suffix=""):
+    """
+    Plot the attention mask. 
+    Attentino mask needs to be constructed using reconstruct_full_mask()
+    """
+    if mask is not None:
+        # Reconstruct full mask from block structure
+        full_mask = reconstruct_full_mask(mask)
+        mask_2d = full_mask.cpu().numpy()
+
+        # Create binary attention pattern plot
+        plt.figure(figsize=(10, 10))
+        # Show sample or full mask depending on size
+        display_mask = mask_2d[:sample_size, :sample_size] if mask_2d.shape[0] > sample_size else mask_2d
+        plt.imshow(display_mask, cmap='Blues', aspect='equal')
+        plt.xlabel('Key Position')
+        plt.ylabel('Query Position')
+        plt.title(f'Attention Mask')
+        plt.colorbar(label='Attention Allowed')
+
+        # Generate filename with timestamp and config values
+        # timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        save_path = f"figures/attention_mask/mask_{title_suffix}.png"
+
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        plt.close()
+        print(f"Attention mask saved to {save_path}")
+
+def plot_random_samples_in_grid(data, 
+                                num_samples=100, 
+                                grid_rows=10, 
+                                grid_cols=10, 
+                                save_path='figures/enc_out_samples_grid.png', 
+                                title='100 Random Samples from encoder output'):
+    """
+    Plot 100 random samples from xxx in a 10x10 grid and save as PNG.
+    """
+    random_indices = torch.randperm(data.shape[0])[:num_samples].cpu().numpy()
+
+    fig, axes = plt.subplots(grid_rows, grid_cols, figsize=(20, 20))
+    fig.suptitle(title, fontsize=16)
+
+    for idx, ax in enumerate(axes.flat):
+        sample_idx = random_indices[idx]
+        sample = data[sample_idx, :].float().detach().cpu().numpy()
+        ax.plot(sample)
+        ax.set_title(f'S{sample_idx}', fontsize=6)
+        ax.tick_params(labelsize=4)
+        ax.grid(True, alpha=0.3)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=150, bbox_inches='tight')
+    plt.close()

xattn.py

@@ -0,0 +1,431 @@
+from .lingua_transformer import *
+
+
+def apply_rotary_emb_xattn(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    seq_dim: int,
+    freqs_cis_q: torch.Tensor,
+    freqs_cis_k: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    xq_ = xq.reshape(*xq.shape[:-1], -1, 1, 2)  # B S H D -> B S H D/2 1 2
+    xk_ = xk.reshape(*xk.shape[:-1], -1, 1, 2)  # B S H D -> B S H D/2 1 2
+    freqs_cis_q = reshape_for_broadcast(
+        freqs_cis_q, xq_, seq_dim
+    ).float()  # S D/2 2 2 -> 1 S 1 D/2 2 2
+    freqs_cis_k = reshape_for_broadcast(
+        freqs_cis_k, xk_, seq_dim
+    ).float()  # S D/2 2 2 -> 1 S 1 D/2 2 2
+    xq_out = (xq_ * freqs_cis_q).sum(5).flatten(3)
+    xk_out = (xk_ * freqs_cis_k).sum(5).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class AdaRMSNorm(nn.Module):
+    """
+    Initialize the RMSNorm normalization layer.
+
+    Args:
+        dim (int): The dimension of the input tensor.
+        eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+    Attributes:
+        eps (float): A small value added to the denominator for numerical stability.
+        weight (nn.Parameter): Learnable scaling parameter.
+
+    """
+
+    def __init__(self, emb_dim, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Linear(emb_dim, dim, bias=True)
+
+    def _norm(self, x: torch.Tensor):
+        return x * torch.rsqrt((x * x).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: torch.Tensor, c: torch.Tensor):
+        x = probe.log_stats(x, "resid")
+        output = self._norm(x.float())
+        return (output * self.weight(c).float()).type_as(x)
+
+    def reset_parameters(self):
+        # bias to ones, weight to 0s
+        nn.init.ones_(self.weight.bias)
+        nn.init.zeros_(self.weight.weight)
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        head_dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        rope_theta: float,
+        rope_dim: int,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.head_dim = head_dim
+        self.rope_theta = rope_theta
+        self.rope_dim = rope_dim
+
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.heads_per_group = self.n_heads // self.n_kv_heads
+
+        self.wq = nn.Linear(
+            dim,
+            n_heads * head_dim,
+            bias=False,
+        )
+        self.wk = nn.Linear(
+            dim,
+            n_kv_heads * head_dim,
+            bias=False,
+        )
+        self.wv = nn.Linear(
+            dim,
+            n_kv_heads * head_dim,
+            bias=False,
+        )
+
+        self.wo = nn.Linear(
+            n_heads * head_dim,
+            dim,
+            bias=False,
+        )
+
+    def forward(
+        self,
+        xq: torch.Tensor,
+        xkv: torch.Tensor,
+        freq_cis: torch.Tensor,
+        tok_idx: Optional[torch.Tensor] = None,
+        cross_tok_idx: Optional[torch.Tensor] = None,
+        mask: Optional[Union[BlockMask, str]] = None,
+        attn_impl: str = "sdpa",
+    ) -> torch.Tensor:
+        # B S D
+        assert attn_impl == "flex_attention", "Only flex_attention is supported for now"
+        bsz, seq_len_q, dim = xq.shape
+        _, seq_len_kv, _ = xkv.shape
+        xq = self.wq(xq.view_as(xq))
+        xk = self.wk(xkv.view_as(xkv))
+        xv = self.wv(xkv.view_as(xkv))
+
+        output_shape = xq.shape
+        # B S D -> B S H D
+        xq = xq.view(bsz, seq_len_q, self.n_heads, self.head_dim)
+        xk = xk.view(bsz, seq_len_kv, self.n_kv_heads, self.head_dim)
+        xv = xv.view(bsz, seq_len_kv, self.n_kv_heads, self.head_dim)
+
+        if self.rope_dim==0:
+            pass
+        elif self.rope_dim==1:
+            if tok_idx is not None and cross_tok_idx is not None:
+                xq, xk = apply_rotary_emb_xattn(
+                    xq, xk, 1, freq_cis[tok_idx], freq_cis[cross_tok_idx]
+                )
+            else:
+                xq, xk = apply_rotary_emb_xattn(
+                    xq, xk, 1, freq_cis[0:seq_len_q], freq_cis[0:seq_len_kv]
+                )
+        elif self.rope_dim==4:
+
+            # Build freqcis_4RoPE by indexing freq_cis with each dimension of tok_idx separately and concatenating
+            # Cat along a new dimension to get [S, head_dim//2, 2, 2]
+            freqcis_parts = []
+            freqcis_cross_parts = []
+            for i in range(self.rope_dim):
+                freqcis_parts.append(freq_cis[tok_idx[:, i]])
+                freqcis_cross_parts.append(freq_cis[cross_tok_idx[:, i]])
+            freqcis_4RoPE = torch.cat(freqcis_parts, dim=1)
+            freqcis_cross_4RoPE = torch.cat(freqcis_cross_parts, dim=1)
+
+
+            xq, xk = apply_rotary_emb_xattn(
+                xq, xk, 1, freqcis_4RoPE, freqcis_cross_4RoPE
+            )
+
+        else:
+            print(f"I dont know how to handle {self.rope_dim=} inside xattn.CrossAttention.forward")
+            import IPython; print('\n\nDebug:'); IPython.embed(); import time;  time.sleep(0.3)
+
+        # This condition helps us be easily compatible
+        # with inference by adding a pluggable KVCache
+        if hasattr(self, "kv_cache"):
+            xk, xv = self.kv_cache.update(xk, xv, tok_idx)
+
+        xk = repeat_kv(xk, self.heads_per_group, dim=2)
+        xv = repeat_kv(xv, self.heads_per_group, dim=2)
+
+        if attn_impl == "flex_attention":
+            assert mask is None or isinstance(mask, BlockMask)
+            xq, xk, xv = map(lambda e: e.transpose(1, 2), (xq, xk, xv))
+            if xq.device.type == "mps":
+                # MPS does not support flex_attention; fall back to SDPA with dense mask
+                if mask is not None:
+                    S_q, S_kv = xq.shape[2], xk.shape[2]
+                    q_idx = torch.arange(S_q, device='cpu')
+                    kv_idx = torch.arange(S_kv, device='cpu')
+                    dense_bool = mask.mask_mod(0, 0, q_idx.unsqueeze(1), kv_idx.unsqueeze(0))
+                    attn_mask = torch.zeros(1, 1, S_q, S_kv, dtype=xq.dtype, device=xq.device)
+                    attn_mask.masked_fill_(~dense_bool.unsqueeze(0).unsqueeze(0).to(xq.device), float("-inf"))
+                else:
+                    attn_mask = None
+                output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask)
+            elif xq.device.type == "cuda":
+                output = flex_attention_comp(xq, xk, xv, block_mask=mask)
+            else:
+                output = flex_attention(xq, xk, xv, block_mask=mask)
+            output = output.transpose(1, 2).contiguous()  # B H S D -> B S H D
+
+        elif attn_impl == "sdpa":
+            xq, xk, xv = map(lambda e: e.transpose(1, 2), (xq, xk, xv))
+            assert mask is None or isinstance(mask, (str, torch.Tensor))
+            is_causal = (mask == "causal") if isinstance(mask, str) else False
+            mask = mask if isinstance(mask, torch.Tensor) else None
+            output = F.scaled_dot_product_attention(
+                xq,
+                xk,
+                xv,
+                is_causal=is_causal,
+                attn_mask=mask,
+            )
+            output = output.transpose(1, 2).contiguous()  # B H S D -> B S H D
+        else:
+            raise NotImplementedError(
+                f"Attention implementation {attn_impl} not supported"
+            )
+
+        output = self.wo(output.reshape(output_shape))
+
+        return output
+
+    def reset_parameters(self, init_std=None, factor=1.0):
+        init_std = init_std or (self.dim ** (-0.5))
+
+        for w in [self.wq, self.wk, self.wv]:
+            nn.init.trunc_normal_(
+                w.weight,
+                mean=0.0,
+                std=init_std,
+                a=-3 * init_std,
+                b=3 * init_std,
+            )
+
+        nn.init.trunc_normal_(
+            self.wo.weight,
+            mean=0.0,
+            std=init_std / factor,
+            a=-3 * init_std,
+            b=3 * init_std,
+        )
+
+
+class FourierConditioner(nn.Module):
+    def __init__(
+        self,
+        output_dim: int,
+        input_dim: int = 1,
+        std: float = 0.02,
+        min_val: float = 0.0,
+        max_val: float = 1.0,
+    ):
+        super().__init__()
+        assert input_dim == 1
+        assert output_dim % 2 == 0
+        self.output_dim = output_dim
+        self.register_buffer("weight", torch.randn([output_dim // 2, input_dim]) * std)
+        self.min_val, self.max_val = min_val, max_val
+        self.proj = nn.Linear(output_dim, output_dim)
+
+    def forward(self, x: list[float], device=None):
+        x = (x - self.min_val) / (self.max_val - self.min_val)
+        f = (2 * torch.pi * x.float() @ self.weight.T).type_as(x)
+        return self.proj(torch.cat([f.cos(), f.sin()], dim=-1))
+
+    def reset_parameters(self, init_std=None, factor=1.0):
+        init_std = init_std or (self.output_dim ** (-0.5))
+
+        self.register_buffer("weight", torch.randn([self.output_dim // 2, 1]).to(self.proj.weight.device) * init_std)
+
+        nn.init.trunc_normal_(
+            self.proj.weight,
+            mean=0.0,
+            std=init_std / factor,
+            a=-3 * init_std,
+            b=3 * init_std,
+        )
+        nn.init.zeros_(self.proj.bias)
+
+
+@dataclass
+class DecoderArgs(BaseTransformerArgs):
+
+    t_dim: int = 64
+    n_heads: int = 8
+    seqlen_t: bool = False
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, args: DecoderArgs):
+        super().__init__()
+
+        assert (args.head_dim is not None) or (args.n_heads is not None), (
+            "Should specify at least head_dim or n_heads"
+        )
+        self.head_dim = args.head_dim or args.dim // args.n_heads
+        self.n_heads = args.n_heads or args.dim // args.head_dim
+        self.n_kv_heads = args.n_kv_heads or self.n_heads
+
+        assert args.n_heads % self.n_kv_heads == 0
+        assert args.dim % args.n_heads == 0
+
+        self.cross_attention = CrossAttention(
+            dim=args.dim,
+            head_dim=self.head_dim,
+            n_heads=self.n_heads,
+            n_kv_heads=self.n_kv_heads,
+            rope_theta=args.rope_theta,
+            rope_dim=args.rope_dim,
+        )
+        self.cross_attention_x_norm = AdaRMSNorm(
+            args.t_dim, args.dim, eps=args.norm_eps
+        )
+
+        self.seqlen_t = args.seqlen_t
+        if args.seqlen_t:
+            self.cross_attention_y_norm = RMSNorm(
+                args.dim, eps=args.norm_eps
+            )
+        else:
+            self.cross_attention_y_norm = AdaRMSNorm(
+                args.t_dim, args.dim, eps=args.norm_eps
+            )
+
+        self.attention = Attention(
+            dim=args.dim,
+            head_dim=self.head_dim,
+            n_heads=self.n_heads,
+            n_kv_heads=self.n_kv_heads,
+            rope_theta=args.rope_theta,
+            rope_dim=args.rope_dim,
+        )
+        self.feed_forward = FeedForward(
+            dim=args.dim,
+            hidden_dim=4 * args.dim,
+            multiple_of=args.multiple_of,
+            ffn_dim_multiplier=args.ffn_dim_multiplier,
+        )
+        self.attention_norm = AdaRMSNorm(args.t_dim, args.dim, eps=args.norm_eps)
+        self.ffn_norm = AdaRMSNorm(args.t_dim, args.dim, eps=args.norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        c: torch.Tensor,
+        freq_cis: torch.Tensor,
+        tok_idx: Optional[torch.Tensor] = None,
+        cross_tok_idx: Optional[torch.Tensor] = None,
+        self_attn_mask: Optional[Union[BlockMask, str]] = None,
+        cross_attn_mask: Optional[Union[BlockMask, str]] = None,
+        attn_impl: str = "sdpa",
+        do_idx: Optional[torch.Tensor] = None,
+        print_layerwise_activation_stats: bool = False,
+    ) -> torch.Tensor:
+        
+        if print_layerwise_activation_stats and do_idx is not None:
+
+            x_normed = self.cross_attention_x_norm(x, c)
+            y_normed = self.cross_attention_y_norm(y, c) if not self.seqlen_t else self.cross_attention_y_norm(y) 
+
+            print(f"\n\tDecoder cross_attn_x_norm: (drop-out) mean={x[:, do_idx, :].mean().item():.6f}, std={x[:, do_idx, :].std().item():.6f}", end=" --> ")
+            print(f"mean={x_normed[:, do_idx, :].mean().item():.6f}, std={x_normed[:, do_idx, :].std().item():.6f}") 
+                        
+            print(f"\tDecoder cross_attn_x_norm: (non-drop) mean={x[:, ~do_idx, :].mean().item():.6f}, std={x[:, ~do_idx, :].std().item():.6f}", end=" --> ")
+            print(f"mean={x_normed[:, ~do_idx, :].mean().item():.6f}, std={x_normed[:, ~do_idx, :].std().item():.6f}") 
+
+            print(f"\n\tDecoder cross_attn_y_norm: (drop-out) mean={y[:, do_idx, :].mean().item():.6f}, std={y[:, do_idx, :].std().item():.6f}", end=" --> ") 
+            print(f"mean={y_normed[:, do_idx, :].mean().item():.6f}, std={y_normed[:, do_idx, :].std().item():.6f}")
+
+            print(f"\tDecoder cross_attn_y_norm: (non-drop) mean={y[:, ~do_idx, :].mean().item():.6f}, std={y[:, ~do_idx, :].std().item():.6f}", end=" --> ")
+            print(f"mean={y_normed[:, ~do_idx, :].mean().item():.6f}, std={y_normed[:, ~do_idx, :].std().item():.6f}") 
+        
+            x = x + self.cross_attention( 
+                x_normed, 
+                y_normed, 
+                freq_cis,
+                tok_idx=tok_idx,
+                cross_tok_idx=cross_tok_idx,
+                mask=cross_attn_mask,
+                attn_impl=attn_impl,
+            )
+
+        else:
+            x = x + self.cross_attention( 
+                self.cross_attention_x_norm(x, c),
+                self.cross_attention_y_norm(y, c) if not self.seqlen_t else self.cross_attention_y_norm(y),
+                freq_cis,
+                tok_idx=tok_idx,
+                cross_tok_idx=cross_tok_idx,
+                mask=cross_attn_mask,
+                attn_impl=attn_impl,
+            )
+
+
+
+        if print_layerwise_activation_stats and do_idx is not None:
+
+            x_normed = self.attention_norm(x, c) 
+
+            print(f"\n\tDecoder self attn_norm: (drop-out) mean={x[:, do_idx, :].mean().item():.6f}, std={x[:, do_idx, :].std().item():.6f}", end=" --> ") 
+            print(f" mean={x_normed[:, do_idx, :].mean().item():.6f}, std={x_normed[:, do_idx, :].std().item():.6f}") 
+            
+            print(f"\tDecoder self attn_norm: (non-drop) mean={x[:, ~do_idx, :].mean().item():.6f}, std={x[:, ~do_idx, :].std().item():.6f}", end=" --> ") 
+            print(f"mean={x_normed[:, ~do_idx, :].mean().item():.6f}, std={x_normed[:, ~do_idx, :].std().item():.6f}") 
+        
+            h = x + self.attention( 
+                x_normed, 
+                freq_cis,
+                tok_idx=tok_idx,
+                mask=self_attn_mask,
+                attn_impl=attn_impl,
+            )
+
+            h_normed = self.ffn_norm(h, c) 
+
+            print(f"\n\tDecoder ffn_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"\tDecoder ffn_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}") 
+
+            out = h + self.feed_forward(h_normed) 
+
+        else:
+
+            h = x + self.attention(
+                self.attention_norm(x, c),
+                freq_cis,
+                tok_idx=tok_idx,
+                mask=self_attn_mask,
+                attn_impl=attn_impl,
+            )
+            out = h + self.feed_forward(self.ffn_norm(h, c))
+
+        return out
+
+    def init_weights(self, init_std=None, factor=1.0):
+        self.cross_attention.reset_parameters(init_std, factor)
+        self.cross_attention_x_norm.reset_parameters()
+        self.cross_attention_y_norm.reset_parameters()
+
+        self.attention.reset_parameters(init_std, factor)
+        self.attention_norm.reset_parameters()
+
+        self.feed_forward.reset_parameters(init_std, factor)
+        self.ffn_norm.reset_parameters()