Initial upload: GUIDO-small 200M bugfix ckpt + Vathos + modeling + README

README.md

@@ -0,0 +1,152 @@
+---
+license: other
+library_name: pytorch
+tags:
+- math
+- code
+- pretraining
+- custom-architecture
+- vathos
+- pico
+language:
+- en
+pipeline_tag: text-generation
+---
+
+# GUIDO-small 200M (bugfixed)
+
+A 195M-parameter math/code-oriented language model, pretrained from scratch on a math-heavy corpus.
+Custom architecture (Vathos PiCOFormer): gated attention, smear-gate, identity-at-init, UDLP MLP,
+RoPE, tied embeddings, fp32-master weights with bf16 autocast.
+
+> ⚠️ This is a **research checkpoint**. Architecture is custom (not HF-compatible) — requires the
+> bundled `vathos/` package + `modeling_guido.py` to load.
+
+## Eval
+
+| | greedy | maj@16 (T=0.6, top_p=0.95, top_k=40) | pass@16 |
+|---|---|---|---|
+| **GSM8K** | 43.0% | **57.5%** | 77.5% |
+| **MATH-500** | 21.0% | **30.0%** | 53.75% |
+
+Numbers from `eval_distill.py` (greedy, N=200/100) and `eval_sc.py` (maj@16, N=80/80).
+
+## Architecture
+
+- **Shape**: 24 layers × 768 d_model × 12 heads × 64 head_dim × 3072 d_ff
+- **Vocab**: 32,768 (Mathstral SPM tokenizer)
+- **Position**: RoPE base 1e6, max_len 8192
+- **Norm**: RMSNorm, qk-norm
+- **Attention**: Multihead Gated (Modded-NanoGPT PR#117 style sparse output gate, gate_input_dim=12)
+- **MLP**: VariableUDLP with LeakyReLU² activation, identity-at-init
+- **Logit softcap**: 30.0
+- **Tied embedding**: yes
+- **Total params**: 195,092,992 (~195M)
+
+## Training
+
+- **Tokens seen**: 4.19B (16000 iter × 4 GPUs × 32 batch × 2048 seq)
+- **Optimizer**: Muon (Newton-Schulz5) on 2D matrices (lr=0.01) + Adam on embedding (lr=0.005) + scalars (lr=0.04)
+- **Schedule**: WSD (warmup 300 → plateau → linear warmdown last 60% to lr×1e-3)
+- **Precision**: fp32 master weights + bf16 autocast (cce kernel gets bf16 cast at boundary)
+- **Hardware**: 4× A100 64GB (Leonardo CINECA), ~3h wall
+
+### Data mixture (corpus_v2)
+
+| weight | shard | tokens (B) |
+|---|---|---|
+| 0.34 | OpenMathInstruct-2 (14M docs) | 6.79 |
+| 0.14 | TinyGSM (full) | 2.70 |
+| 0.10 | OpenMathReasoning subset | 1.50 |
+| 0.07 | NuminaMath-1.5 | 0.41 |
+| 0.05 | NuminaMath-CoT | 0.48 |
+| 0.22 | FineWeb-Edu score≥3 (sample-10BT) | 3.00 |
+| 0.08 | Cosmopedia subset | 1.19 |
+
+Math 70% / NL 30%. Tokenizer: `mistralai/Mathstral-7B-v0.1`. Decontamination: canonical-hash vs MATH-500 and GSM8K test (note: not fuzzy-decon).
+
+## How to load
+
+### Quick start
+
+```bash
+pip install -r requirements.txt
+```
+
+```python
+from huggingface_hub import snapshot_download
+local_dir = snapshot_download(repo_id="Paerle/GUIDO_test_200M", token="<your_token>")
+
+import sys, torch
+sys.path.insert(0, local_dir + "/vathos")
+sys.path.insert(0, local_dir)
+import modeling_guido as m
+from transformers import AutoTokenizer
+
+# Inferenza arch dal ckpt (così non serve config)
+state = torch.load(local_dir + "/pytorch_model.bin", map_location="cpu", weights_only=False)
+sd = state["model"]
+hp = m.HP()
+hp.head_dim = 64
+hp.d_model = sd["backbone.embedder.embedding.weight"].shape[1]
+hp.n_heads = hp.d_model // hp.head_dim
+hp.d_ff = sd["backbone.blocks.0.channel_mixer.expand.weight"].shape[0]
+hp.n_layers = max(int(k.split("blocks.")[1].split(".")[0]) for k in sd if "blocks." in k) + 1
+model = m.PiCOFormerLM(hp, vocab_size=sd["backbone.embedder.embedding.weight"].shape[0])
+model.load_state_dict(sd, strict=False)
+model = model.cuda().bfloat16().eval()
+tok = AutoTokenizer.from_pretrained("mistralai/Mathstral-7B-v0.1")
+```
+
+See [`inference_example.py`](./inference_example.py) for a complete generation example
+(temperature, top-p, top-k sampling).
+
+## Files
+
+| file | description |
+|---|---|
+| `pytorch_model.bin` | model weights (fp32 master, 745MB) |
+| `config.json` | architecture + HP + eval results |
+| `modeling_guido.py` | model definition (`PiCOFormerLM` + Muon optimizer + reader; only PiCOFormerLM/HP needed for inference) |
+| `inference_example.py` | minimal load + generate example |
+| `vathos/` | bundled Vathos package (custom backbone — see github.com/MarioPaerle/Aplos) |
+| `requirements.txt` | deps (torch, transformers, cut-cross-entropy) |
+
+## Known limitations
+
+- **Decontamination**: canonical-hash only. ~7% of MATH-500 has near-duplicates in OpenMathInstruct-2
+  augmentation (audited). Real MATH-500 capability is ~21% greedy, not inflated by overt memorization.
+- **No instruction tuning**: this is a base pretrain. No SFT / RLHF.
+- **Domain**: heavily math-skewed. NL is 30% (FineWeb-Edu + Cosmopedia) — basic fluency but not chat-style.
+- **Tokenizer dependency**: requires `mistralai/Mathstral-7B-v0.1` tokenizer from HF.
+
+## Hardware/precision notes
+
+- Inference works fine in **bf16** (just call `.bfloat16()`). Weights are stored fp32 — 745MB on disk
+  but ~390MB GPU memory after `.bfloat16()`.
+- Custom CUDA kernels: uses `cut_cross_entropy` (Apple) for fused linear+CE in training.
+  For inference you can replace with standard `F.linear + F.cross_entropy`.
+
+## Architecture sketch
+
+```
+    input_ids
+        │
+   embedder (tied) ───────────��──────────────┐
+        │                                    │
+   smear_gate (sparse on first 12 dims) ─────│ x0
+        │                                    │
+   [24× block]                                │
+    ├── spatial: gated_attn + RoPE + qk-norm │
+    ├── channel: VariableUDLP (LeakyReLU²)    │
+    └── + x0_lambda × x0  ───────────────────┘
+        │
+   final_norm
+        │
+   tied unembed → cce / softcap 30 → logits
+```
+
+## License & citation
+
+Vathos backbone (custom by Mario Paerle): https://github.com/MarioPaerle/Aplos
+Model weights: research/private. Contact owner for usage terms.

config.json

@@ -0,0 +1,44 @@
+{
+  "model_type": "guido_small",
+  "model_name": "GUIDO-small 200M (bugfixed)",
+  "architecture": "Vathos PiCOFormer (custom)",
+  "version": "v1-bugfix",
+  "n_layers": 24,
+  "d_model": 768,
+  "n_heads": 12,
+  "head_dim": 64,
+  "d_ff": 3072,
+  "vocab_size": 32768,
+  "max_seq_len": 8192,
+  "rope_base": 1000000.0,
+  "logit_softcap": 30.0,
+  "tied_embeddings": true,
+  "qk_norm": true,
+  "use_gated_attn": true,
+  "use_smear_gate": true,
+  "gate_input_dim": 12,
+  "mlp_kind": "udlp",
+  "orthogonalizer": "ns5",
+  "muon_backend_steps": 5,
+  "training_dtype": "fp32-master/bf16-autocast",
+  "checkpoint_dtype": "float32",
+  "training_steps": 16000,
+  "training_tokens": 4194304000,
+  "training_data": "corpus_v2 mixture 70/22/8 (math/fineweb/cosmopedia)",
+  "training_mixture": {
+    "openmath_full": 0.34,
+    "tinygsm": 0.14,
+    "openmathreasoning": 0.10,
+    "numina_15": 0.07,
+    "numina_cot": 0.05,
+    "fineweb_edu": 0.22,
+    "cosmopedia": 0.08
+  },
+  "tokenizer": "mistralai/Mathstral-7B-v0.1",
+  "params_total": 195092992,
+  "params_total_human": "195M",
+  "eval_gsm8k_greedy_pct": 43.0,
+  "eval_gsm8k_maj16_pct": 57.5,
+  "eval_math500_greedy_pct": 21.0,
+  "eval_math500_maj16_pct": 30.0
+}

inference_example.py

@@ -0,0 +1,94 @@
+"""Load GUIDO-small 200M (bugfix) e genera testo.
+Run:  python inference_example.py
+"""
+import sys
+from pathlib import Path
+import torch
+from transformers import AutoTokenizer
+
+# Aggiungi questa cartella al sys.path così Vathos viene trovato
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE / "vathos"))      # bundled Vathos package
+sys.path.insert(0, str(HERE))                  # modeling_guido.py
+
+# Il modeling è in modeling_guido.py (fork di train_guido_small.py — usa la classe PiCOFormerLM).
+# `import modeling_guido as m` registra Vathos in production mode + carica le classi.
+import modeling_guido as m
+
+
+def load_model(ckpt_path: str, device="cuda"):
+    """Carica il ckpt fp32, infera l'architettura dai pesi, restituisce model+tokenizer."""
+    state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    sd = state["model"] if "model" in state else state
+    # strip eventuali prefissi
+    sd = {k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k: v for k, v in sd.items()}
+
+    # Infer arch dal state_dict (così funziona anche se il config.json non c'è)
+    hp = m.HP()
+    hp.head_dim = 64
+    hp.d_model = sd["backbone.embedder.embedding.weight"].shape[1]
+    hp.n_heads = hp.d_model // hp.head_dim
+    hp.d_ff = sd["backbone.blocks.0.channel_mixer.expand.weight"].shape[0]
+    hp.n_layers = max(int(k.split("blocks.")[1].split(".")[0])
+                      for k in sd if "blocks." in k) + 1
+    vocab = sd["backbone.embedder.embedding.weight"].shape[0]
+    print(f"[load] arch inferred: {hp.n_layers}L × {hp.d_model} × {hp.n_heads}h × {hp.d_ff}ff  vocab={vocab}")
+
+    model = m.PiCOFormerLM(hp, vocab_size=vocab).to(device).bfloat16()
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    if missing:    print(f"[load] missing keys ({len(missing)}): {missing[:3]}...")
+    if unexpected: print(f"[load] unexpected ({len(unexpected)}): {unexpected[:3]}...")
+    model.eval()
+
+    tok = AutoTokenizer.from_pretrained("mistralai/Mathstral-7B-v0.1")
+    return model, tok
+
+
+@torch.no_grad()
+def generate(model, tok, prompt: str, max_new=200, temperature=0.6, top_p=0.95, top_k=40, device="cuda"):
+    """Sampling semplice (T, top-p, top-k)."""
+    import torch.nn.functional as F
+    ids = tok(prompt, return_tensors="pt").input_ids.to(device)
+    eos = tok.eos_token_id or 2
+    softcap = model.logit_softcap
+    for _ in range(max_new):
+        h = model._hidden(ids)
+        logits = F.linear(h[:, -1, :], model.classifier_weight)
+        if softcap > 0:
+            logits = softcap * torch.tanh(logits / softcap)
+        logits = logits.float() / temperature
+        if top_k and top_k < logits.size(-1):
+            kth = logits.topk(top_k, dim=-1).values[:, -1:]
+            logits = logits.masked_fill(logits < kth, float("-inf"))
+        probs = F.softmax(logits, dim=-1)
+        if top_p < 1.0:
+            sp, si = probs.sort(dim=-1, descending=True)
+            csum = sp.cumsum(dim=-1)
+            mask = csum - sp > top_p
+            sp = sp.masked_fill(mask, 0.0); sp /= sp.sum(dim=-1, keepdim=True)
+            choice = torch.multinomial(sp, 1)
+            nxt = si.gather(-1, choice)
+        else:
+            nxt = torch.multinomial(probs, 1)
+        ids = torch.cat([ids, nxt], dim=1)
+        if nxt.item() == eos:
+            break
+    return tok.decode(ids[0], skip_special_tokens=False)
+
+
+if __name__ == "__main__":
+    # Adatta il path al checkpoint
+    CKPT = HERE / "pytorch_model.bin"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model, tok = load_model(str(CKPT), device=device)
+
+    prompts = [
+        "Question: Sarah has 5 apples and gives 2 to Tom. How many does she have left?\nAnswer: ",
+        "Solve: $x^2 + 3x + 2 = 0$. Step 1:",
+        "The largest planet in the solar system is",
+    ]
+    for p in prompts:
+        print("=" * 60)
+        print("PROMPT:", p)
+        out = generate(model, tok, p, max_new=120, device=device)
+        print("OUT:   ", out)

modeling_guido.py

@@ -0,0 +1,760 @@
+"""Guido-small — pretrain ~200M (24L×768), math-first chattabile.
+
+Fork di `train_distill_start.py` SENZA KD (il teacher non aiuta a questa scala —
+vedi SESSION_HANDOFF.md §A). Differenze chiave:
+  - shape 200M (24L × 768 × 12h × 64hd × 3072ff)
+  - MultiShardMixture reader: legge corpus_v2 (7 shard per-dataset, NON shufflati a
+    write-time) con chunk-shuffle a read-time (NO rimpiazzo, copertura 100%, ordine
+    casuale). Pesi mixture 70% math / 22% fineweb / 8% cosmopedia in HP.mixture.
+  - loss-trace per-step → CSV (loss, gnorm, lr, math_frac) con UN sync ogni
+    `train_log_every` step (buffer su GPU). math_frac correla spike↔batch math-heavy.
+  - looping opzionale (env LOOP_STYLE / LOOP_ACTIVATION_FRAC), default OFF per la baseline.
+
+Run:
+    torchrun --standalone --nproc_per_node=4 train_guido_small.py
+
+Pre-req: corpus_v2 tokenizzato (Mathstral SPM 32k, uint32) in $FAST/corpus_v2/<name>/.
+"""
+from __future__ import annotations
+
+import csv
+import os
+import re
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+# ---- Vathos backbone ----
+APLOS_PATH = os.environ.get("APLOS_PATH", "/leonardo_work/IscrC_YENDRI/paerle/PiCO/aplos")
+if APLOS_PATH not in sys.path:
+    sys.path.insert(0, APLOS_PATH)
+from Vathos._basics import (Builder, RMSNorm as VRMSNorm, ReLU2, LeakyReLU2,
+                            VariableUDLP, VariableGatedUDLP, set_vathos_mode)
+from Vathos._spatials import MultiheadAttentionMixer, MultiheadGatedAttentionMixer, RoPE
+from Vathos.blocks import PiCOFormer as VathosPiCOFormer, SmearGate
+set_vathos_mode("production")
+
+# FAST RMSNorm fused
+def _fast_rmsnorm_forward(self, x):
+    return F.rms_norm(x, (x.size(-1),), self.weight, self.eps)
+VRMSNorm.forward = _fast_rmsnorm_forward
+
+# ---- Shard reader ----
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "PiCO2"))
+from pico2.shards import ShardReader
+
+from cut_cross_entropy import linear_cross_entropy
+
+
+# =============================================================================
+# HYPERPARAMETERS
+# =============================================================================
+class HP:
+    # data — corpus_v2 multi-shard mixture (Mathstral SPM 32k)
+    corpus_root = os.environ.get(
+        "CORPUS_ROOT",
+        "/leonardo_scratch/fast/IscrC_YENDRI/mprignan/corpus_v2",
+    )
+    # Mixture preset selezionabile via MIXTURE_PRESET env. Default = 70/22/8 (run baseline 200M).
+    # "v3_noloops" = fineweb 22%→15% (≈2/3), delta freed → math (77/15/8 total).
+    _MIXTURE_PRESETS = {
+        "default": {
+            "openmath_full": 0.34, "tinygsm": 0.14, "openmathreasoning": 0.10,
+            "numina_15": 0.07, "numina_cot": 0.05,          # = 0.70 math
+            "fineweb_edu": 0.22, "cosmopedia": 0.08,        # = 0.30 NL
+        },
+        "v3_noloops": {
+            "openmath_full": 0.37, "tinygsm": 0.15, "openmathreasoning": 0.11,
+            "numina_15": 0.08, "numina_cot": 0.06,          # = 0.77 math
+            "fineweb_edu": 0.15, "cosmopedia": 0.08,        # = 0.23 NL (fineweb ridotto a 2/3)
+        },
+    }
+    mixture = _MIXTURE_PRESETS[os.environ.get("MIXTURE_PRESET", "default")]
+    nl_datasets = ("fineweb_edu", "cosmopedia")          # per math_frac logging
+    seq_len = 2048
+    bs_per_dev = int(os.environ.get("BS_PER_DEV", 32))
+    train_batch_tokens = 4 * bs_per_dev * 2048            # 4×32×2048 = 262k tok/step (4 GPU)
+
+    # model — shape 200M (scale-up dal 92M: d_model 512→768, d_ff 2048→3072, heads 8→12)
+    n_layers = int(os.environ.get("N_LAYERS", 24))
+    d_model = int(os.environ.get("D_MODEL", 768))
+    n_heads = int(os.environ.get("N_HEADS", 12))
+    head_dim = 64
+    d_ff = int(os.environ.get("D_FF", 3072))
+    rope_base = 1_000_000.0
+    logit_softcap = 30.0
+    tied_embeddings = True
+    qk_norm = True
+    # v3 features (locked)
+    use_gated_attn = True
+    # Channel mixer (FFN) variant. Knob aperto dall'audit A/B 2026-05-28 (vedi analysis/guido_ab/REPORT).
+    #   "udlp"          : VariableUDLP, contract(LeakyReLU²(expand(x)))                 (current default)
+    #   "gated_udlp"    : VariableGatedUDLP, UDLP·sigmoid(gate_proj(x[..., :K]))         (sparse output gate)
+    #   "leaky_reglu2"  : LeakyReGLU² GLU, contract(LeakyReLU²(expand(x)) · up(x))       (act·value, +1 proj)
+    mlp_kind = os.environ.get("MLP_KIND", "udlp")
+    use_smear_gate = True
+    # gate_input_dim: 12 era stato scelto stile parameter-golf/sparse. Sul d=768 è ~1.5% di x:
+    # per il prossimo run testare 64 (8%) o 128 (17%) — vedi audit gate-activity (40% dei gate<0.1).
+    gate_input_dim = int(os.environ.get("GATE_INPUT_DIM", 12))
+
+    # optimizer (Muon LOCKED)
+    optimizer_kind = "muon"
+    # Orthogonalizer per Muon. Knob aperto dall'audit A/B 2026-05-28.
+    #   "ns5"          : Newton-Schulz quintic, coeff (3.4445,-4.7750,2.0315), Keller Jordan
+    #   "polar_express": Polar Express (arxiv 2505.16932) — schedule di 8 coeff/step,
+    #                    convergenza più rapida. RACCOMANDATO MUON_BACKEND_STEPS=8 (default 5
+    #                    di NS5 è sotto-iterato per PE → convergenza parziale).
+    orthogonalizer = os.environ.get("ORTHOGONALIZER", "ns5")
+    muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5))
+    matrix_lr = 0.01
+    tied_embed_lr = 0.005
+    scalar_lr = 0.04
+    muon_momentum = 0.95
+    muon_momentum_warmup_start = 0.85
+    muon_momentum_warmup_steps = 200
+    beta1 = 0.9
+    beta2 = 0.95
+    adam_eps = 1e-8
+    grad_clip = 1.0
+
+    # schedule (WSD: warmup → plateau → linear warmdown a lr_min_scale)
+    iterations = int(os.environ.get("ITERATIONS", 15000))   # ~4B tok; CALIBRA dopo lo smoke
+    warmup_steps = int(os.environ.get("WARMUP", 300))
+    warmdown_iters = int(os.environ.get("WARMDOWN", int(0.60 * iterations)))  # 60%
+    lr_min_scale = float(os.environ.get("LR_MIN_SCALE", 0.001))
+
+    # === Layer looping (opzionale, default OFF per baseline) ===
+    # loop_v2 sul 92M: +4.5pp GSM8K @ activation_frac 0.30 (durante plateau LR).
+    loop_style = os.environ.get("LOOP_STYLE", "")
+    loop_activation_frac = float(os.environ.get("LOOP_ACTIVATION_FRAC", 0.30))
+    loop_pre_warm = os.environ.get("LOOP_PRE_WARM", "1") == "1"
+
+    # logging / checkpoint
+    train_log_every = int(os.environ.get("LOG_EVERY", 50))   # dump CSV ogni N (loss PER-STEP nel buffer)
+    loss_trace_csv = os.environ.get(
+        "LOSS_TRACE_CSV",
+        "/leonardo_work/IscrC_YENDRI/paerle/PiCO/Guido-1/PiCO2_test/logs/guido_small_loss_trace.csv",
+    )
+    ckpt_dir = os.environ.get(
+        "CKPT_DIR",
+        "/leonardo_work/IscrC_YENDRI/paerle/PiCO/ckpts/pico_guido_small",
+    )
+    save_final = True
+    save_interval = int(os.environ.get("SAVE_INTERVAL", 2000))
+    keep_n_recent = 5
+    seed = int(os.environ.get("SEED", 1337))
+
+
+# =============================================================================
+# MUON OPTIMIZER
+# =============================================================================
+def zeropower_via_newtonschulz5(G: Tensor, steps: int = 5, eps: float = 1e-7) -> Tensor:
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G.bfloat16()
+    X /= X.norm() + eps
+    transposed = G.size(0) > G.size(1)
+    if transposed:
+        X = X.T
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * A @ A
+        X = a * X + B @ X
+    return X.T if transposed else X
+
+
+# Polar Express coefficients (arxiv 2505.16932). Stessa forma di NS5 ma SCHEDULE di 8 set:
+# iter 1 aggressivo per matrici mal condizionate, iter 7-8 = steady-state Halley (3/8, -10/8, 15/8).
+# Source: NVIDIA NeMo Emerging-Optimizers (authoritative, production-tested).
+# https://docs.nvidia.com/nemo/emerging-optimizers/0.1.0/_modules/emerging_optimizers/orthogonalized_optimizers/muon_utils.html
+_POLAR_EXPRESS_COEFFS = (
+    (8.2051, -22.9019, 16.4607),
+    (4.0664,  -2.8612,  0.5184),
+    (3.9096,  -2.8234,  0.5250),
+    (3.2856,  -2.4153,  0.4853),
+    (2.2779,  -1.6198,  0.3985),
+    (1.8726,  -1.2307,  0.3585),
+    (1.8564,  -1.2132,  0.3568),
+    (1.8750,  -1.2500,  0.3750),
+)
+
+
+def zeropower_via_polar_express(G: Tensor, steps: int = 8, eps: float = 1e-7) -> Tensor:
+    """Polar Express orthogonalizer — convergenza più rapida di NS5 quintic.
+    Iteration: X = a*X + (b*A + c*A@A)@X  con A=X@X^T, a/b/c SCHEDULATI per step.
+    Normalizzazione Frobenius (identica a NS5). Per convergenza piena servono ≥7 step
+    (raccomandato MUON_BACKEND_STEPS=8). Oltre 8 step usa lo steady-state (Halley)."""
+    X = G.bfloat16()
+    X /= X.norm() + eps
+    transposed = G.size(0) > G.size(1)
+    if transposed:
+        X = X.T
+    n_coeffs = len(_POLAR_EXPRESS_COEFFS)
+    for k in range(steps):
+        a, b, c = _POLAR_EXPRESS_COEFFS[k if k < n_coeffs else n_coeffs - 1]
+        A = X @ X.T
+        B = b * A + c * A @ A
+        X = a * X + B @ X
+    return X.T if transposed else X
+
+
+class Muon(torch.optim.Optimizer):
+    def __init__(self, params, lr, momentum, backend_steps, nesterov=True):
+        super().__init__(params, dict(lr=lr, momentum=momentum,
+                                      backend_steps=backend_steps, nesterov=nesterov))
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        distributed = dist.is_available() and dist.is_initialized()
+        world_size = dist.get_world_size() if distributed else 1
+        rank = dist.get_rank() if distributed else 0
+        for group in self.param_groups:
+            params = group["params"]
+            if not params:
+                continue
+            lr, mom, ns_steps, nesterov = group["lr"], group["momentum"], group["backend_steps"], group["nesterov"]
+            total = sum(int(p.numel()) for p in params)
+            updates_flat = torch.zeros(total, device=params[0].device, dtype=torch.bfloat16)
+            curr = 0
+            for i, p in enumerate(params):
+                if i % world_size == rank and p.grad is not None:
+                    g = p.grad
+                    state = self.state[p]
+                    if "momentum_buffer" not in state:
+                        state["momentum_buffer"] = torch.zeros_like(g)
+                    buf = state["momentum_buffer"]
+                    buf.mul_(mom).add_(g)
+                    if nesterov:
+                        g = g.add(buf, alpha=mom)
+                    g = zeropower_via_newtonschulz5(g, steps=ns_steps)
+                    g *= max(1, g.size(0) / g.size(1)) ** 0.5
+                    updates_flat[curr:curr + p.numel()] = g.reshape(-1)
+                curr += p.numel()
+            if distributed:
+                dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM)
+            curr = 0
+            for p in params:
+                u = updates_flat[curr:curr + p.numel()].view_as(p).to(p.dtype)
+                p.add_(u, alpha=-lr)
+                curr += p.numel()
+
+
+# =============================================================================
+# CHANNEL MIXER VARIANT — LeakyReGLU² GLU-style FFN
+# =============================================================================
+class LeakyReGLU2_FFN(nn.Module):
+    """LeakyReLU² GLU: contract(LeakyReLU²(expand(x)) · up(x)).
+
+    GLU-style channel mixer (à la SwiGLU/ReGLU) con activation LeakyReLU². Aggiunge UN extra
+    proiezione `up` (d_model→M) rispetto a VariableUDLP → +50% param nell'FFN. Identity-at-init:
+    contract=0 ⇒ branch nullo all'init (gradiente fluisce attraverso `expand` e `up`).
+
+    Compatibile con Vathos `Builder`: signature (d_model, d_output, M, activation, dropout).
+    Param naming: `expand` (act path, ndim==2 → Muon), `up` (value path), `contract` (output).
+    """
+    def __init__(self, d_model, d_output, M, activation=None, dropout=0.0):
+        super().__init__()
+        self.expand   = nn.Linear(d_model, M, bias=False)   # act path (the W in act(xW))
+        self.up       = nn.Linear(d_model, M, bias=False)   # value path (the V in ·xV)
+        self.contract = nn.Linear(M, d_output, bias=False)
+        self.activation = activation() if activation is not None else LeakyReLU2()
+        self.dropout = nn.Dropout(dropout)
+        self._init_weights()
+
+    def _init_weights(self):
+        torch.nn.init.orthogonal_(self.expand.weight)
+        torch.nn.init.orthogonal_(self.up.weight)
+        torch.nn.init.zeros_(self.contract.weight)   # identity-at-init: branch nullo
+
+    def forward(self, x):
+        return self.dropout(self.contract(self.activation(self.expand(x)) * self.up(x)))
+
+
+# =============================================================================
+# DATA — MultiShardMixture: chunk-shuffle a read-time (NO rimpiazzo)
+# =============================================================================
+class ShuffledStream:
+    """Un dataset → finestre (seq_len+1) in ordine PERMUTATO (shuffle senza rimpiazzo).
+    Chunk disgiunti per rank; re-permuta a fine epoca (sub-epoca = nessun wrap col budget attuale)."""
+    def __init__(self, shard_dir, rank, world_size, seq_len, seed):
+        self.reader = ShardReader(shard_dir)
+        self.win = seq_len + 1
+        self.n_chunks = self.reader.total_tokens // self.win
+        self.rank, self.ws, self.seed, self._epoch = rank, world_size, seed, 0
+        self._reperm()
+
+    def _reperm(self):
+        rng = np.random.default_rng(self.seed + 104729 * self._epoch)
+        self.my = rng.permutation(self.n_chunks)[self.rank::self.ws]
+        self.ptr = 0
+        self._epoch += 1
+
+    def next_window(self):
+        if self.ptr >= len(self.my):
+            self._reperm()
+        ci = int(self.my[self.ptr])
+        self.ptr += 1
+        return self.reader.read(ci * self.win, self.win)   # np.uint32 [win]
+
+
+class MultiShardMixture:
+    """Weighted-pick del dataset per-sample + finestra chunk-shuffled. Ogni doc ≤1 volta/epoca."""
+    def __init__(self, mixture, corpus_root, rank, world_size, device, seq_len, seed):
+        self.names = list(mixture.keys())
+        self.w = np.array([mixture[n] for n in self.names], float)
+        self.w /= self.w.sum()
+        self.streams = {
+            n: ShuffledStream(str(Path(corpus_root) / n), rank, world_size, seq_len, seed + i * 7919)
+            for i, n in enumerate(self.names)
+        }
+        self.rng = np.random.default_rng(seed * 31 + rank)
+        self.device = device
+
+    def next_batch(self, bs):
+        picks = self.rng.choice(len(self.names), size=bs, p=self.w)
+        xs, ys, ids = [], [], []
+        for pi in picks:
+            w = self.streams[self.names[pi]].next_window()
+            t = torch.from_numpy(w.astype(np.int64))
+            xs.append(t[:-1])
+            ys.append(t[1:])
+            ids.append(self.names[pi])
+        return (torch.stack(xs).to(self.device, non_blocking=True),
+                torch.stack(ys).to(self.device, non_blocking=True), ids)
+
+
+# =============================================================================
+# LAYER LOOPING — parser + scalar gate init=0 (Universal Transformer style)
+# =============================================================================
+_LOOP_SEG_RE = re.compile(r"\(\s*([\d\s,]+?)\s*\)\s*x\s*(\d+)")
+
+
+def parse_loop_style(s: str, n_layers: int):
+    """Parse loop style stringa → lista di (indices, n_total_executions).
+    Esempio: "(3,4,5)x2, (8,9)x3" → [([3,4,5], 2), ([8,9], 3)]. Semantica xN: N esecuzioni
+    totali (1 originale + N-1 loop extra con gate). Validation: in-range, contigui, disgiunti, N>=2."""
+    s = s.strip()
+    if not s:
+        return []
+    stripped = re.sub(r"\s+", "", s)
+    pattern_strict = re.compile(r"\(([\d,]+)\)x(\d+)")
+    consumed = "".join(f"({a})x{b}" for a, b in pattern_strict.findall(stripped))
+    consumed_with_commas = ",".join(f"({a})x{b}" for a, b in pattern_strict.findall(stripped))
+    if stripped != consumed and stripped != consumed_with_commas:
+        raise ValueError(f"loop_style mal formato: {s!r}. Atteso: '(i,j,k)xN, (l,m)xM, ...'")
+
+    segments = []
+    used_indices = set()
+    for indices_str, n_str in _LOOP_SEG_RE.findall(s):
+        indices = sorted({int(x.strip()) for x in indices_str.split(",") if x.strip()})
+        n = int(n_str)
+        if n < 2:
+            raise ValueError(f"loop count xN deve avere N>=2 (segmento ({indices}) ha x{n}). "
+                             f"Usa loop_style='' per nessun looping.")
+        if not indices:
+            raise ValueError(f"loop segment vuoto in {s!r}")
+        if indices != list(range(min(indices), max(indices) + 1)):
+            raise ValueError(f"segmento {tuple(indices)} NON contiguo: i layer in un loop "
+                             f"devono essere contigui (e.g. (3,4,5), non (3,5)).")
+        for i in indices:
+            if not 0 <= i < n_layers:
+                raise ValueError(f"layer {i} fuori range [0, {n_layers}) nel segmento {tuple(indices)}")
+            if i in used_indices:
+                raise ValueError(f"layer {i} appare in più segmenti — i segmenti devono essere disgiunti")
+            used_indices.add(i)
+        segments.append((indices, n))
+    segments.sort(key=lambda seg: seg[0][0])
+    for s1, s2 in zip(segments, segments[1:]):
+        if s1[0][-1] >= s2[0][0]:
+            raise ValueError(f"segmenti sovrapposti dopo ordinamento: {s1[0]} e {s2[0]}")
+    return segments
+
+
+class LoopGate(nn.Module):
+    """Scalar gate init=0 (additive): h_main + α·(h_post − h_main). α=0 → identity."""
+    def __init__(self):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.zeros(()))
+
+    def forward(self, h_main: Tensor, h_post: Tensor) -> Tensor:
+        return h_main + self.alpha * (h_post - h_main)
+
+
+# =============================================================================
+# MODEL — wrapper Vathos + layer looping (CE only, no KD)
+# =============================================================================
+class PiCOFormerLM(nn.Module):
+    def __init__(self, hp: HP, vocab_size: int, loop_style: str = ""):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.logit_softcap = hp.logit_softcap
+        self.tied_embeddings = hp.tied_embeddings
+        self.loop_style = loop_style
+        self.loop_segments = parse_loop_style(loop_style, hp.n_layers) if loop_style else []
+        self.loop_gates = nn.ModuleList([
+            nn.ModuleList([LoopGate() for _ in range(n - 1)])
+            for indices, n in self.loop_segments
+        ])
+        self._seg_starts = {indices[0]: seg_id for seg_id, (indices, _) in enumerate(self.loop_segments)}
+        self.looped_mode = False
+        rope = RoPE(dim=hp.head_dim, max_len=8192, base=hp.rope_base)
+        if hp.use_gated_attn:
+            attn_builder = Builder(MultiheadGatedAttentionMixer,
+                n_heads=hp.n_heads, causal=True, dropout=0.0,
+                qk_norm=hp.qk_norm, pos_emb=rope, gate_input_dim=hp.gate_input_dim)
+        else:
+            attn_builder = Builder(MultiheadAttentionMixer,
+                n_heads=hp.n_heads, causal=True, dropout=0.0,
+                qk_norm=hp.qk_norm, pos_emb=rope)
+        if hp.mlp_kind == "udlp":
+            chan_builder = Builder(VariableUDLP,
+                d_output=hp.d_model, M=hp.d_ff, activation=LeakyReLU2)
+        elif hp.mlp_kind == "gated_udlp":
+            chan_builder = Builder(VariableGatedUDLP,
+                d_output=hp.d_model, M=hp.d_ff, activation=LeakyReLU2,
+                gate_input_dim=hp.gate_input_dim)
+        elif hp.mlp_kind == "leaky_reglu2":
+            chan_builder = Builder(LeakyReGLU2_FFN,
+                d_output=hp.d_model, M=hp.d_ff, activation=LeakyReLU2)
+        else:
+            raise ValueError(f"unknown mlp_kind={hp.mlp_kind!r}; usa udlp|gated_udlp|leaky_reglu2")
+        smear = SmearGate(hp.d_model, gate_input_dim=hp.gate_input_dim) if hp.use_smear_gate else None
+        smear_lookback = 1 if hp.use_smear_gate else 0
+        self.backbone = VathosPiCOFormer(
+            vocab_size=vocab_size, d_model=hp.d_model, n_layers=hp.n_layers,
+            spatials=attn_builder, channel=chan_builder, norm=VRMSNorm,
+            ve_groups=None, smear_gate=smear, smear_gate_lookback=smear_lookback,
+            logit_softcap=hp.logit_softcap, tied_embeddings=hp.tied_embeddings,
+        )
+
+    @property
+    def classifier_weight(self):
+        if self.tied_embeddings:
+            return self.backbone.embedder.embedding.weight
+        return self.backbone.unembedder.linear.weight
+
+    def _hidden(self, input_ids: Tensor) -> Tensor:
+        if self.looped_mode and self.loop_segments:
+            return self._hidden_looped(input_ids)
+        return self._hidden_normal(input_ids)
+
+    def _hidden_normal(self, input_ids: Tensor) -> Tensor:
+        bb = self.backbone
+        x0 = bb.embedder(input_ids)
+        if getattr(bb, "smear_gate", None) is not None:
+            x0 = bb.smear_gate(x0)
+        ves = bb._compute_ves(input_ids) if hasattr(bb, "_compute_ves") else [None] * len(bb.blocks)
+        h = x0
+        for i, block in enumerate(bb.blocks):
+            h = block(h, x0, ve=ves[i] if ves is not None else None)
+        h = bb.final_norm(h)
+        return h
+
+    def _hidden_looped(self, input_ids: Tensor) -> Tensor:
+        bb = self.backbone
+        x0 = bb.embedder(input_ids)
+        if getattr(bb, "smear_gate", None) is not None:
+            x0 = bb.smear_gate(x0)
+        ves = bb._compute_ves(input_ids) if hasattr(bb, "_compute_ves") else [None] * len(bb.blocks)
+        h = x0
+        n_layers = len(bb.blocks)
+        i = 0
+        while i < n_layers:
+            seg_id = self._seg_starts.get(i)
+            if seg_id is None:
+                h = bb.blocks[i](h, x0, ve=ves[i] if ves is not None else None)
+                i += 1
+                continue
+            indices, n_total = self.loop_segments[seg_id]
+            for li in indices:
+                h = bb.blocks[li](h, x0, ve=ves[li] if ves is not None else None)
+            h_main = h
+            for k in range(n_total - 1):
+                h_post = h_main
+                for li in indices:
+                    h_post = bb.blocks[li](h_post, x0, ve=ves[li] if ves is not None else None)
+                h_main = self.loop_gates[seg_id][k](h_main, h_post)
+            h = h_main
+            i = indices[-1] + 1
+        h = bb.final_norm(h)
+        return h
+
+    def forward(self, input_ids: Tensor, targets: Tensor) -> Tensor:
+        """CE only (cce fused, no logits materialized)."""
+        h = self._hidden(input_ids)
+        W = self.classifier_weight
+        # cce vuole input bf16/fp16; con master fp32 + autocast h/W possono essere fp32 → cast esplicito
+        # (input cce identici a prima; il grad risale al master fp32 attraverso il cast).
+        return linear_cross_entropy(
+            h.reshape(-1, h.size(-1)).bfloat16(), W.bfloat16(), targets.reshape(-1),
+            reduction="mean",
+            softcap=self.logit_softcap if self.logit_softcap > 0 else None,
+        )
+
+
+# =============================================================================
+# MAIN
+# =============================================================================
+def main():
+    hp = HP()
+
+    # Dispatch orthogonalizer (Muon.step chiama il nome globale `zeropower_via_newtonschulz5`,
+    # quindi se l'utente sceglie polar_express ri-bindo quel nome alla nuova fn).
+    global zeropower_via_newtonschulz5
+    if hp.orthogonalizer == "ns5":
+        zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5)
+    elif hp.orthogonalizer == "polar_express":
+        # NotImplementedError lo lancia al primo call dentro Muon.step.
+        zeropower_via_newtonschulz5 = torch.compile(zeropower_via_polar_express)
+    else:
+        raise ValueError(f"unknown orthogonalizer={hp.orthogonalizer!r}; usa ns5|polar_express")
+
+    distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ
+    rank = int(os.environ.get("RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    device = torch.device("cuda", local_rank)
+    torch.cuda.set_device(device)
+    if distributed:
+        dist.init_process_group(backend="nccl", device_id=device)
+        dist.barrier()
+    is_main = (rank == 0)
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cuda.enable_flash_sdp(True)
+    torch.backends.cuda.enable_mem_efficient_sdp(False)
+    torch.backends.cuda.enable_math_sdp(False)
+    torch.backends.cuda.enable_cudnn_sdp(False)
+    import torch._inductor.config as _ind
+    _ind.coordinate_descent_tuning = True
+    _ind.fx_graph_cache = True
+    _ind.triton.cudagraphs = False
+
+    torch.manual_seed(hp.seed)
+    torch.cuda.manual_seed_all(hp.seed)
+    np.random.seed(hp.seed)
+
+    def log(msg):
+        if is_main:
+            print(msg, flush=True)
+
+    log(f"world_size={world_size} rank={rank} device={device}")
+    log(f"corpus_root={hp.corpus_root}")
+    log(f"mixture={hp.mixture}")
+
+    # ---- data: verifica vocab consistency su tutti gli shard ----
+    vocab_size = None
+    eos_id = None
+    for name in hp.mixture:
+        idx = ShardReader(str(Path(hp.corpus_root) / name)).index
+        if vocab_size is None:
+            vocab_size, eos_id = idx.vocab_size, idx.eos_id
+        assert idx.vocab_size == vocab_size, (
+            f"vocab mismatch: {name} ha {idx.vocab_size}, atteso {vocab_size}")
+        assert idx.eos_id == eos_id, f"eos mismatch su {name}"
+    log(f"vocab={vocab_size} eos={eos_id} (consistente su {len(hp.mixture)} shard)")
+    assert vocab_size > eos_id
+
+    train_loader = MultiShardMixture(hp.mixture, hp.corpus_root, rank, world_size,
+                                     device, hp.seq_len, hp.seed)
+    if is_main:
+        for n, s in train_loader.streams.items():
+            log(f"  [stream] {n:<20} chunks={s.n_chunks:,} (per-rank {len(s.my):,}) w={hp.mixture[n]:.2f}")
+
+    # ---- student model ----
+    # FIX (audit bf16): pesi in fp32 = MASTER weights. Il forward gira comunque in bf16 via autocast.
+    # Senza master fp32, al floor del warmdown gli update (~lr·1/√fan) cadono sotto la ULP bf16 e
+    # p.add_ diventa un no-op → il warmdown non abbassa più la loss. fp32 master lo risolve.
+    model = PiCOFormerLM(hp, vocab_size=vocab_size, loop_style=hp.loop_style).to(device)
+    n_params = sum(p.numel() for p in model.parameters())
+    log(f"model_params={n_params/1e6:.2f}M  shape={hp.n_layers}L×{hp.d_model}×{hp.n_heads}h×{hp.d_ff}ff  dtype=fp32-master/bf16-autocast")
+    log(f"knobs: mlp_kind={hp.mlp_kind}  gate_input_dim={hp.gate_input_dim}  orthogonalizer={hp.orthogonalizer}  muon_steps={hp.muon_backend_steps}  mixture_preset={os.environ.get('MIXTURE_PRESET','default')}")
+    if hp.loop_style:
+        loop_summary = ", ".join(f"({','.join(map(str, idx))})x{n}" for idx, n in model.loop_segments)
+        n_extra_passes = sum(len(idx) * (n - 1) for idx, n in model.loop_segments)
+        log(f"[loop] style={hp.loop_style!r}  segments={loop_summary}")
+        log(f"[loop] gates={sum(len(g) for g in model.loop_gates)} (init α=0)  "
+            f"extra_passes_per_fwd={n_extra_passes}  activation_frac={hp.loop_activation_frac}")
+    compiled = torch.compile(model, dynamic=False, fullgraph=False, mode="max-autotune-no-cudagraphs")
+    ddp_kwargs = dict(device_ids=[local_rank], broadcast_buffers=False, gradient_as_bucket_view=True)
+    if hp.loop_style:
+        ddp_kwargs["find_unused_parameters"] = True   # loop_gates senza grad in NORMAL fwd
+    model_for_train = DDP(compiled, **ddp_kwargs) if distributed else compiled
+    model.train()
+
+    # ---- optimizer ----
+    block_params = list(model.backbone.blocks.named_parameters())
+    matrix_params = [p for n, p in block_params if p.ndim == 2]
+    scalar_params = [p for n, p in block_params if p.ndim < 2]
+    for n, p in model.backbone.final_norm.named_parameters():
+        scalar_params.append(p)
+    for n, p in model.loop_gates.named_parameters():
+        scalar_params.append(p)
+    # FIX (audit): smear_gate è attributo top-level del backbone (NON dentro blocks) → senza questo
+    # i suoi param cadono fuori da ogni optimizer e restano congelati a init (smear_lambda=0 → gate
+    # no-op per tutto il run). Routing per ndim, coerente col resto (gate.weight 2D→Muon, lambda→scalar).
+    if getattr(model.backbone, "smear_gate", None) is not None:
+        for n, p in model.backbone.smear_gate.named_parameters():
+            (matrix_params if p.ndim == 2 else scalar_params).append(p)
+    embed_param = model.backbone.embedder.embedding.weight
+
+    opt_muon = Muon(matrix_params, lr=hp.matrix_lr, momentum=hp.muon_momentum,
+                    backend_steps=hp.muon_backend_steps)
+    opt_embed = torch.optim.Adam(
+        [{"params": [embed_param], "lr": hp.tied_embed_lr, "base_lr": hp.tied_embed_lr}],
+        betas=(hp.beta1, hp.beta2), eps=hp.adam_eps, fused=True,
+    )
+    opt_scalar = torch.optim.Adam(
+        [{"params": scalar_params, "lr": hp.scalar_lr, "base_lr": hp.scalar_lr}],
+        betas=(hp.beta1, hp.beta2), eps=hp.adam_eps, fused=True,
+    )
+    optimizers = [opt_muon, opt_embed, opt_scalar]
+    for g in opt_muon.param_groups:
+        g["base_lr"] = hp.matrix_lr
+    log(f"opt: Muon={sum(p.numel() for p in matrix_params)/1e6:.2f}M "
+        f"Adam embed={embed_param.numel()/1e6:.2f}M scalars={sum(p.numel() for p in scalar_params)/1e6:.2f}M")
+
+    local_bs = hp.bs_per_dev
+    log(f"bs/rank={local_bs} global_bs={local_bs*world_size} tok/step={local_bs*world_size*hp.seq_len:,}")
+    log(f"iterations={hp.iterations} warmup={hp.warmup_steps} warmdown={hp.warmdown_iters} "
+        f"→ target_tok={hp.iterations*local_bs*world_size*hp.seq_len/1e9:.2f}B")
+
+    def lr_scale(step):
+        if step < hp.warmup_steps:
+            return step / max(hp.warmup_steps, 1)
+        decay_start = hp.iterations - hp.warmdown_iters
+        if step < decay_start:
+            return 1.0
+        progress = (step - decay_start) / max(hp.warmdown_iters, 1)
+        return max(1.0 - progress * (1.0 - hp.lr_min_scale), hp.lr_min_scale)
+
+    Path(hp.ckpt_dir).mkdir(parents=True, exist_ok=True)
+
+    # ---- loss-trace CSV (per-step; UN sync ogni train_log_every) ----
+    csv_f = csv_w = None
+    if is_main:
+        Path(hp.loss_trace_csv).parent.mkdir(parents=True, exist_ok=True)
+        csv_f = open(hp.loss_trace_csv, "w", newline="")
+        csv_w = csv.writer(csv_f)
+        csv_w.writerow(["step", "loss", "gnorm", "lr", "math_frac"])
+        log(f"loss_trace_csv={hp.loss_trace_csv}")
+    loss_buf, gn_buf, lr_buf, step_buf, mf_buf = [], [], [], [], []
+
+    # ---- Layer looping: pre-compile both variants ----
+    loop_activation_step = int(hp.iterations * hp.loop_activation_frac) if hp.loop_style else None
+    loop_active = False
+    if hp.loop_style and hp.loop_pre_warm:
+        log(f"[loop] pre-compiling NORMAL + LOOPED graphs (one-shot, può richiedere 5-15 min)")
+        x_warm = torch.zeros((local_bs, hp.seq_len), dtype=torch.long, device=device)
+        y_warm = torch.zeros_like(x_warm)
+        with torch.no_grad():
+            model.looped_mode = False
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+                _ = compiled(x_warm, y_warm)
+            torch.cuda.synchronize()
+            log(f"[loop] NORMAL fwd graph compiled")
+            model.looped_mode = True
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+                _ = compiled(x_warm, y_warm)
+            torch.cuda.synchronize()
+            log(f"[loop] LOOPED fwd graph compiled")
+        model.looped_mode = False
+        del x_warm, y_warm
+
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    total_tok_seen = 0
+    nan_abort = False
+
+    for step in range(1, hp.iterations + 1):
+        if loop_activation_step is not None and not loop_active and step >= loop_activation_step:
+            log(f"[loop] ACTIVATING looping at step {step} (gates α=0, identity passthrough)")
+            model.looped_mode = True
+            loop_active = True
+        for opt in optimizers:
+            opt.zero_grad(set_to_none=True)
+
+        x, y, ds_ids = train_loader.next_batch(local_bs)
+        total_tok_seen += local_bs * world_size * hp.seq_len
+
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
+            loss = model_for_train(x, y)
+
+        loss.backward()
+
+        scale = lr_scale(step)
+        for opt in optimizers:
+            for g in opt.param_groups:
+                g["lr"] = g["base_lr"] * scale
+        frac = min(step / hp.muon_momentum_warmup_steps, 1.0) if hp.muon_momentum_warmup_steps > 0 else 1.0
+        mom = (1 - frac) * hp.muon_momentum_warmup_start + frac * hp.muon_momentum
+        for g in opt_muon.param_groups:
+            g["momentum"] = mom
+
+        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), hp.grad_clip)
+        for opt in optimizers:
+            opt.step()
+
+        # per-step buffers (NO sync qui — un solo sync ogni train_log_every al flush)
+        loss_buf.append(loss.detach())
+        gn_buf.append(gn.detach())
+        lr_buf.append(opt_muon.param_groups[0]["lr"])
+        step_buf.append(step)
+        mf_buf.append(sum(d not in hp.nl_datasets for d in ds_ids) / len(ds_ids))
+
+        if step % hp.train_log_every == 0 or step == 1 or step == hp.iterations:
+            torch.cuda.synchronize()
+            ls = torch.stack(loss_buf).float().cpu().tolist()
+            gs = torch.stack(gn_buf).float().cpu().tolist()
+            if not all(np.isfinite(ls)):
+                log(f"!! NaN/Inf loss intorno a step {step}, abort")
+                nan_abort = True
+            if is_main:
+                for s, l, g, lr_, mf in zip(step_buf, ls, gs, lr_buf, mf_buf):
+                    csv_w.writerow([s, f"{l:.5f}", f"{g:.3f}", f"{lr_:.3e}", f"{mf:.2f}"])
+                csv_f.flush()
+                elapsed = time.perf_counter() - t0
+                tps = total_tok_seen / max(elapsed, 1e-9)
+                loop_tag = ""
+                if loop_active:
+                    alphas = [g.alpha.detach().abs().item()
+                              for seg_gates in model.loop_gates for g in seg_gates]
+                    loop_tag = f" [LOOP |α|avg={sum(alphas)/max(len(alphas),1):.3f}]"
+                log(f"step {step:>5d}/{hp.iterations}{loop_tag} loss={ls[-1]:.4f} "
+                    f"gnorm={gs[-1]:.3f} lr={lr_buf[-1]:.2e} mathfrac={mf_buf[-1]:.2f} "
+                    f"tok={total_tok_seen/1e9:.2f}B tok/s={tps:>10,.0f} elapsed={elapsed:.1f}s")
+            loss_buf.clear(); gn_buf.clear(); lr_buf.clear(); step_buf.clear(); mf_buf.clear()
+
+        if nan_abort:
+            break
+
+        if (hp.save_interval > 0 and is_main and step % hp.save_interval == 0
+                and step != hp.iterations):
+            ckpt_path = Path(hp.ckpt_dir) / f"step_{step:06d}.pt"
+            torch.save({"step": step, "model": model.state_dict(), "vocab_size": vocab_size}, ckpt_path)
+            log(f"  [ckpt] saved {ckpt_path.name}")
+            intermediates = [p for p in sorted(Path(hp.ckpt_dir).glob("step_*.pt")) if "_final" not in p.name]
+            for old in intermediates[:-hp.keep_n_recent]:
+                old.unlink()
+
+    if hp.save_final and is_main and not nan_abort:
+        ckpt_path = Path(hp.ckpt_dir) / f"step_{hp.iterations:06d}_final.pt"
+        torch.save({"step": hp.iterations, "model": model.state_dict(), "vocab_size": vocab_size}, ckpt_path)
+        log(f"saved {ckpt_path}")
+    if is_main and csv_f is not None:
+        csv_f.close()
+
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

pytorch_model.bin

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

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.4
+numpy
+transformers>=4.40
+cut-cross-entropy
+huggingface_hub>=1.0

vathos/Analyzer.py

@@ -0,0 +1,706 @@
+"""
+Vathos ModdedFormer Analyzer API
+=================================
+A pure-Python/Jupyter alternative to the Streamlit inspector.
+Designed for rigorous, programmatic topological and statistical inspection
+of Transformer parameters, gradients, and activation manifolds.
+
+Usage in a Jupyter Notebook:
+----------------------------
+    import torch
+    from vathos_analyzer import VathosAnalyzer
+
+    # 1. Initialize
+    analyzer = VathosAnalyzer(model)
+
+    # 2. Plot Architecture Overview & Parameter distributions
+    analyzer.plot_overview()
+    analyzer.plot_weight_distributions(layer_idx=0, component="both")
+
+    # 3. Manifold Extraction via Forward Hooks
+    x = torch.randint(0, model.vocab_size, (1, 64))
+    analyzer.capture_forward(x) # Installs hooks, runs forward, saves state
+
+    # 4. Explore Layer Topology (Attention & FFN mapping f(x))
+    analyzer.plot_attention_weights(layer_idx=0)
+    analyzer.plot_ffn_manifold(layer_idx=0)
+
+    # 5. Training History
+    analyzer.plot_training_history()
+
+    # --- OR USE THE ONE-LINER MACRO ---
+    # from vathos_analyzer import study
+    # analyzer = study(model)
+"""
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import matplotlib.pyplot as plt
+from typing import Optional, Dict, List, Union
+
+try:
+    from IPython.display import display
+    import pandas as pd
+
+    HAS_IPYTHON = True
+except ImportError:
+    HAS_IPYTHON = False
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  DARK THEME STYLING
+# ──────────────────────────────────────────────────────────────────────────────
+DARK_THEME = {
+    "figure.facecolor": "#0D0D1A",
+    "axes.facecolor": "#0D0D1A",
+    "axes.edgecolor": "#252550",
+    "axes.labelcolor": "#C0C0E0",
+    "xtick.color": "#9CA3AF",
+    "ytick.color": "#9CA3AF",
+    "text.color": "#D8D8F0",
+    "grid.color": "#1E1E3A",
+    "grid.alpha": 0.6,
+    "legend.facecolor": "#11111F",
+    "legend.edgecolor": "#252550",
+}
+COLORS = ["#A78BFA", "#60A5FA", "#34D399", "#F472B6", "#FB923C",
+          "#FBBF24", "#38BDF8", "#A3E635", "#E879F9", "#F87171"]
+
+
+def apply_theme():
+    plt.rcParams.update(DARK_THEME)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  CORE ANALYZER CLASS
+# ──────────────────────────────────────────────────────────────────────────────
+class VathosAnalyzer:
+    def __init__(self, model: nn.Module, dark_mode: bool = True):
+        """
+        Initializes the analyzer.
+        :param model: The Vathos ModdedFormer instance.
+        """
+        self.model = model
+        self.captured_activations: Dict[int, Dict[str, torch.Tensor]] = {}
+        if dark_mode:
+            apply_theme()
+
+    # ─── EXTRACTION LOGIC ─────────────────────────────────────────────────────
+
+    def _is_attention(self, mixer) -> bool:
+        cls_name = type(mixer).__name__.lower()
+        return "attention" in cls_name or "attn" in cls_name
+
+    def _tensor_stats(self, t: torch.Tensor) -> dict:
+        f = t.detach().float().cpu()
+        v = f.numpy().flatten()
+        return {
+            "shape": tuple(t.shape),
+            "numel": t.numel(),
+            "mean": float(v.mean()),
+            "std": float(v.std()),
+            "min": float(v.min()),
+            "max": float(v.max()),
+            "l2": float(np.linalg.norm(v)),
+            "sparsity": float((np.abs(v) < 1e-6).mean()),
+        }
+
+    def _svd_values(self, t: torch.Tensor) -> Optional[np.ndarray]:
+        f = t.detach().float().cpu()
+        if f.ndim == 2:
+            try:
+                return torch.linalg.svdvals(f).numpy()
+            except Exception:
+                return None
+        if f.ndim == 3:
+            results = []
+            for i in range(f.shape[0]):
+                try:
+                    results.append(torch.linalg.svdvals(f[i]).numpy())
+                except Exception:
+                    pass
+            return np.stack(results) if results else None
+        return None
+
+    def _get_attention_weights(self, mixer, x: torch.Tensor) -> Optional[np.ndarray]:
+        if hasattr(mixer, "get_attention_weights"):
+            try:
+                with torch.no_grad():
+                    w = mixer.get_attention_weights(x)
+                return w.detach().float().cpu().numpy()
+            except Exception:
+                pass
+        try:
+            with torch.no_grad():
+                B, L, D = x.shape
+                if hasattr(mixer, "qkv"):
+                    proj = mixer.qkv
+                    n_heads = mixer.n_heads
+                    head_dim = mixer.head_dim
+                    qkv = proj(x).view(B, L, 3, n_heads, head_dim)
+                    q, k, v = qkv.unbind(dim=2)
+                    q, k = q.transpose(1, 2).float(), k.transpose(1, 2).float()
+                    scale = math.sqrt(head_dim)
+                    scores = torch.matmul(q, k.transpose(-2, -1)) / scale
+                    if getattr(mixer, "causal", True):
+                        mask = torch.tril(torch.ones(L, L, device=x.device)).bool()
+                        scores = scores.masked_fill(~mask, float("-inf"))
+                    weights = torch.softmax(scores, dim=-1)
+                    return weights[0].cpu().numpy()
+                elif hasattr(mixer, "qk"):
+                    proj = mixer.qk
+                    n_heads = mixer.n_heads
+                    head_dim = mixer.head_dim
+                    qk = proj(x).view(B, L, 2, n_heads, head_dim)
+                    q, k = qk.unbind(dim=2)
+                    q, k = q.transpose(1, 2).float(), k.transpose(1, 2).float()
+                    scale = math.sqrt(head_dim)
+                    scores = torch.matmul(q, k.transpose(-2, -1)) / scale
+                    if getattr(mixer, "causal", True):
+                        mask = torch.tril(torch.ones(L, L, device=x.device)).bool()
+                        scores = scores.masked_fill(~mask, float("-inf"))
+                    weights = torch.softmax(scores, dim=-1)
+                    return weights[0].cpu().numpy()
+        except Exception:
+            pass
+        return None
+
+    def get_layer_params(self, layer_idx: int, component: str = "both") -> Dict[str, torch.Tensor]:
+        block = self.model.blocks[layer_idx]
+        sm = block.spatial_mixer
+        cm = block.channel_mixer
+
+        selected_modules = {}
+        if component in ["spatial", "both"]:
+            selected_modules["spatial"] = sm
+        if component in ["channel", "both"]:
+            selected_modules["channel"] = cm
+
+        all_params = {}
+        for prefix, module in selected_modules.items():
+            for name, p in module.named_parameters():
+                all_params[f"{prefix}.{name}"] = p
+        return all_params
+
+    # ─── FORWARD PASS HOOKING (MANIFOLD CAPTURE) ──────────────────────────────
+
+    def capture_forward(self, x: torch.Tensor, clear_previous: bool = True):
+        """
+        Runs a forward pass while attaching hooks to capture inputs/outputs of
+        Spatial Mixers, and the strict pre/post activations of Channel Mixers.
+        """
+        if clear_previous:
+            self.captured_activations.clear()
+
+        handles = []
+
+        for i, block in enumerate(self.model.blocks):
+            self.captured_activations[i] = {}
+
+            # Hook 1: Input to Spatial Mixer (For Q/K/V computations)
+            def make_spatial_hook(layer_idx):
+                def hook(m, inp):
+                    self.captured_activations[layer_idx]["spatial_in"] = inp[0].detach()
+
+                return hook
+
+            handles.append(block.spatial_mixer.register_forward_pre_hook(make_spatial_hook(i)))
+
+            cm = block.channel_mixer
+
+            # Hook 2 & 3: Channel Mixer Expansion / Contraction (Pre and Post Activation)
+            if hasattr(cm, 'expand') and hasattr(cm, 'contract'):
+                def make_expand_hook(layer_idx):
+                    def hook(m, inp, out):
+                        self.captured_activations[layer_idx]["channel_pre_act"] = out.detach()
+
+                    return hook
+
+                def make_contract_hook(layer_idx):
+                    def hook(m, inp):
+                        self.captured_activations[layer_idx]["channel_post_act"] = inp[0].detach()
+
+                    return hook
+
+                handles.append(cm.expand.register_forward_hook(make_expand_hook(i)))
+                handles.append(cm.contract.register_forward_pre_hook(make_contract_hook(i)))
+            else:
+                # Fallback per block custom
+                def make_cm_hook(layer_idx):
+                    def hook(m, inp):
+                        self.captured_activations[layer_idx]["channel_in"] = inp[0].detach()
+
+                    return hook
+
+                handles.append(cm.register_forward_pre_hook(make_cm_hook(i)))
+
+        # Run forward pass
+        self.model.eval()
+        try:
+            device = next(self.model.parameters()).device
+        except StopIteration:
+            device = torch.device("cpu")
+
+        with torch.no_grad():
+            self.model(x.to(device))
+
+        # Cleanup hooks
+        for h in handles:
+            h.remove()
+
+        print(f"✅ Manifolds captured for {len(self.model.blocks)} layers.")
+
+    # ─── PLOTTING ROUTINES ────────────────────────────────────────────────────
+
+    def show_fig(self, fig: plt.Figure):
+        """Helper to show a figure inline or wait for script end."""
+        if HAS_IPYTHON:
+            display(fig)
+            plt.close(fig)
+        else:
+            fig.show()
+
+    def plot_overview(self):
+        """Plots topological params: param norms, skip lambdas, zeroskips."""
+        print("=== 🏗️ Model Topological Overview ===")
+        total_params = sum(p.numel() for p in self.model.parameters())
+        print(f"Total Params: {total_params:,} | Layers: {self.model.n_layer} | Embed Dim: {self.model.embed_dim}")
+
+        # 1. Param Norms
+        layer_names, l2_norms = [], []
+        for i, block in enumerate(self.model.blocks):
+            for name, p in block.named_parameters():
+                layer_names.append(f"L{i}.{name}")
+                l2_norms.append(float(p.detach().norm(2)))
+
+        fig, ax = plt.subplots(figsize=(max(8, len(layer_names) * 0.35), 4))
+        ax.bar(range(len(layer_names)), l2_norms, color="#60A5FA", alpha=0.8)
+        ax.set_xticks(range(len(layer_names)))
+        ax.set_xticklabels(layer_names, rotation=90, fontsize=6)
+        ax.set_ylabel("L2 norm")
+        ax.set_title("Per-parameter L2 Norms", color="#A78BFA")
+        ax.grid(True, ls="--", alpha=0.4, axis="y")
+        fig.tight_layout()
+        self.show_fig(fig)
+
+        # 2. Skip Lambdas
+        if hasattr(self.model, "skip_lambdas") and self.model.skip_lambdas:
+            lambdas = {k: float(v.detach()) for k, v in self.model.skip_lambdas.items()}
+            fig2, ax2 = plt.subplots(figsize=(max(6, len(lambdas) * 1.4), 3.5))
+            keys = list(lambdas.keys())
+            vals = [lambdas[k] for k in keys]
+            bars = ax2.bar(range(len(keys)), vals, color=COLORS[:len(keys)], alpha=0.85)
+            ax2.set_xticks(range(len(keys)))
+            ax2.set_xticklabels([k.replace("route_", "").replace("_to_", "→") for k in keys], rotation=30, ha="right",
+                                fontsize=8)
+            ax2.axhline(0, color="#9CA3AF", linewidth=0.8, ls="--")
+            ax2.set_title("Skip-connection Gate Values (λ)", color="#A78BFA")
+            for bar, v in zip(bars, vals):
+                ax2.text(bar.get_x() + bar.get_width() / 2, v + 0.003, f"{v:.4f}", ha="center", va="bottom", fontsize=8,
+                         color="#D8D8F0")
+            fig2.tight_layout()
+            self.show_fig(fig2)
+
+        # 3. Zeroskip Params
+        if hasattr(self.model, "zeroskip") and self.model.zeroskip:
+            vals = [float(p.detach()) for p in self.model.zeroskip_params]
+            fig3, ax3 = plt.subplots(figsize=(max(6, len(vals) * 0.7), 3.5))
+            ax3.plot(vals, marker="o", color="#FBBF24", linewidth=1.5, markersize=6)
+            ax3.axhline(0, color="#9CA3AF", linewidth=0.8, ls="--")
+            ax3.set_xlabel("layer index")
+            ax3.set_ylabel("zeroskip α")
+            ax3.set_title("ZeroSkip Parameters (x₀ coefficient)", color="#A78BFA")
+            ax3.grid(True, ls="--", alpha=0.4)
+            fig3.tight_layout()
+            self.show_fig(fig3)
+
+    def print_layer_stats(self, layer_idx: int, component: str = "both"):
+        """Prints a statistical dataframe of parameters for a given layer."""
+        params = self.get_layer_params(layer_idx, component)
+        rows = []
+        for pname, p in params.items():
+            s = self._tensor_stats(p)
+            rows.append({
+                "Parameter": pname, "Shape": str(s["shape"]), "Numel": s['numel'],
+                "Mean": s['mean'], "Std": s['std'], "Min": s['min'],
+                "Max": s['max'], "L2": s['l2'], "Sparsity": s['sparsity'],
+                "Grad": "Yes" if p.grad is not None else "No",
+            })
+        if HAS_IPYTHON:
+            df = pd.DataFrame(rows)
+            display(df)
+        else:
+            for r in rows:
+                print(r)
+
+    def plot_weight_distributions(self, layer_idx: int, component: str = "both", plot_gradients: bool = False,
+                                  bins: int = 80):
+        params = self.get_layer_params(layer_idx, component)
+        if plot_gradients:
+            params = {k: v.grad for k, v in params.items() if v.grad is not None}
+            if not params:
+                print("No gradients found. Run loss.backward() first.")
+                return
+
+        n = len(params)
+        if n == 0: return
+        cols = min(n, 3)
+        rows = math.ceil(n / cols)
+        fig, axes = plt.subplots(rows, cols, figsize=(5 * cols, 3.5 * rows))
+        if n == 1:
+            axes = np.array([[axes]])
+        elif rows == 1:
+            axes = axes.reshape(1, -1)
+
+        for idx, (name, t) in enumerate(params.items()):
+            ax = axes[idx // cols][idx % cols]
+            v = t.detach().float().cpu().numpy().flatten()
+            st_val = self._tensor_stats(t)
+            ax.hist(v, bins=bins, color=COLORS[idx % len(COLORS)], alpha=0.8, density=True)
+            ax.set_title(f"{name}\nμ={st_val['mean']:.3e}  σ={st_val['std']:.3e}", fontsize=9)
+            ax.grid(True, ls="--", alpha=0.4)
+
+        for idx in range(n, rows * cols):
+            axes[idx // cols][idx % cols].set_visible(False)
+
+        fig.suptitle(f"Layer {layer_idx} - {'Gradient' if plot_gradients else 'Weight'} Distributions", color="#A78BFA",
+                     fontsize=12)
+        fig.tight_layout()
+        self.show_fig(fig)
+
+    def plot_svd_spectrum(self, layer_idx: int, component: str = "both", plot_gradients: bool = False,
+                          log_scale: bool = True):
+        params = self.get_layer_params(layer_idx, component)
+        if plot_gradients:
+            params = {k: v.grad for k, v in params.items() if v.grad is not None}
+
+        eligible = {k: v for k, v in params.items() if v.ndim >= 2}
+        if not eligible:
+            print("No 2D parameters found for SVD.")
+            return
+
+        n = len(eligible)
+        cols = min(n, 3)
+        rows = math.ceil(n / cols)
+        fig, axes = plt.subplots(rows, cols, figsize=(5 * cols, 3.5 * rows))
+        if n == 1:
+            axes = np.array([[axes]])
+        elif rows == 1:
+            axes = axes.reshape(1, -1)
+
+        for idx, (name, t) in enumerate(eligible.items()):
+            ax = axes[idx // cols][idx % cols]
+            sv = self._svd_values(t)
+            if sv is None: continue
+
+            if sv.ndim == 2:
+                for h in range(sv.shape[0]):
+                    ax.plot(sv[h], color=COLORS[idx % len(COLORS)], alpha=0.3, linewidth=0.8)
+                ax.plot(sv.mean(0), color=COLORS[idx % len(COLORS)], linewidth=2, label="mean")
+                ax.legend(fontsize=8)
+            else:
+                ax.plot(sv, color=COLORS[idx % len(COLORS)], linewidth=1.5, marker=".", markersize=3)
+
+            cond = sv.flatten()[0] / (sv.flatten()[-1] + 1e-12)
+            ax.set_title(f"{name} | cond={cond:.1f}", fontsize=9)
+            if log_scale: ax.set_yscale("log")
+            ax.grid(True, ls="--", alpha=0.4)
+
+        for idx in range(n, rows * cols):
+            axes[idx // cols][idx % cols].set_visible(False)
+
+        fig.suptitle(f"Layer {layer_idx} - Singular Value Spectra ($\Sigma$)", color="#A78BFA", fontsize=12)
+        fig.tight_layout()
+        self.show_fig(fig)
+
+    # ─── MANIFOLD EXPLORATION (ATTENTION E FFN) ───────────────────────────────
+
+    def plot_attention_weights(self, layer_idx: int, tokens: Optional[List[str]] = None):
+        """Plots the attention map of the layer using the captured manifold."""
+        if layer_idx not in self.captured_activations or "spatial_in" not in self.captured_activations[layer_idx]:
+            print(f"No captured input for layer {layer_idx}. Run `capture_forward(x)` first.")
+            return
+
+        sm = self.model.blocks[layer_idx].spatial_mixer
+        if not self._is_attention(sm):
+            print(f"Layer {layer_idx} spatial mixer is not Attention (is {type(sm).__name__}).")
+            return
+
+        x = self.captured_activations[layer_idx]["spatial_in"]
+        w = self._get_attention_weights(sm, x)
+        if w is None:
+            print("Could not extract attention weights.")
+            return
+
+        if w.ndim == 2: w = w[np.newaxis]
+        H, L_q, L_k = w.shape
+        cols = min(H, 4)
+        rows = math.ceil(H / cols)
+
+        fig, axes = plt.subplots(rows, cols, figsize=(4 * cols, 3.5 * rows))
+        axes = np.array(axes).flatten() if H > 1 else [axes]
+
+        for h in range(H):
+            ax = axes[h]
+            wh = w[h]
+            im = ax.imshow(wh, aspect="auto", cmap="magma", vmin=0, vmax=wh.max() + 1e-9)
+            ax.set_title(f"Head {h}", fontsize=10, color="#C4B5FD")
+            if tokens and len(tokens) == L_q:
+                clean_tokens = [t.replace('Ġ', ' ') for t in tokens]
+                fs = max(4, min(10, int(300 / L_q)))
+                ax.set_xticks(range(L_k))
+                ax.set_yticks(range(L_q))
+                ax.set_xticklabels(clean_tokens, rotation=90, fontsize=fs, color="#9CA3AF")
+                ax.set_yticklabels(clean_tokens, fontsize=fs, color="#9CA3AF")
+            fig.colorbar(im, ax=ax, shrink=0.8)
+
+        for h in range(H, len(axes)):
+            axes[h].set_visible(False)
+
+        fig.suptitle(f"Layer {layer_idx} - Attention Topology Analysis", color="#A78BFA", fontsize=12)
+        fig.tight_layout()
+        self.show_fig(fig)
+
+    def plot_ffn_manifold(self, layer_idx: int):
+        """
+        Traccia la densità (PDF empirica) delle attivazioni prima della non linearità (XW_exp),
+        sovrapposta alla mappa non-lineare $f(x)$. Questo rivela immediatamente condizioni
+        di annullamento gradiente (dead neurons) o eccessiva saturazione.
+        Dipende dai Forward Hooks eseguiti tramite `capture_forward()`.
+        """
+        if layer_idx not in self.captured_activations:
+            print(f"No manifold captured. Run `capture_forward(x)` first.")
+            return
+
+        cap = self.captured_activations[layer_idx]
+        cm = self.model.blocks[layer_idx].channel_mixer
+
+        if "channel_pre_act" in cap and "channel_post_act" in cap:
+            pre_act = cap["channel_pre_act"][0].cpu().numpy()
+            post_act = cap["channel_post_act"][0].cpu().numpy()
+        elif "channel_in" in cap:
+            # Fallback se non ci sono expand/contract espliciti
+            x = cap["channel_in"]
+            data = self._get_ffn_activations(cm, x, temp=1.0)
+            if not data:
+                print("Could not compute FFN manifold.")
+                return
+            pre_act = data["pre_act"]
+            post_act = data["post_act"]
+        else:
+            print("Missing necessary hooks data for FFN.")
+            return
+
+        # Estrazione act_fn per mapping analitico
+        act_fn = getattr(cm, "activation", getattr(cm, "act", getattr(cm, "act_fn", None)))
+
+        fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
+        pre_flat = pre_act.flatten()
+        post_flat = post_act.flatten()
+
+        # ── Subplot 1: Pre-act vs f(x) ──
+        ax1 = axes[0]
+        ax1_twin = ax1.twinx()
+
+        # Dominio assoluto
+        true_min = float(pre_flat.min())
+        true_max = float(pre_flat.max())
+        x_min = min(-5.0, true_min - 1.0)
+        x_max = max(5.0, true_max + 1.0)
+        x_vals = np.linspace(x_min, x_max, 1000)
+
+        y_vals = None
+        is_empirical = False
+
+        # 1. Tentativo analitico
+        if act_fn is not None:
+            try:
+                device = next(cm.parameters()).device if list(cm.parameters()) else "cpu"
+                dtype = next(cm.parameters()).dtype if list(cm.parameters()) else torch.float32
+                with torch.no_grad():
+                    xt = torch.tensor(x_vals, dtype=dtype, device=device).view(1, 1, -1)
+                    yt = act_fn(xt)
+                    y_vals = yt.cpu().float().numpy().flatten()
+            except Exception:
+                pass
+
+        # 2. Empirico / Fallback
+        if (y_vals is None or np.array_equal(y_vals, x_vals)) and pre_flat.shape == post_flat.shape:
+            if not np.allclose(pre_flat, post_flat):
+                sort_idx = np.argsort(pre_flat)
+                x_vals = pre_flat[sort_idx]
+                y_vals = post_flat[sort_idx]
+                is_empirical = True
+            else:
+                y_vals = x_vals
+        elif y_vals is None:
+            y_vals = x_vals
+
+        label_map = "$f(x)$ map (Empirica)" if is_empirical else "$f(x)$ map"
+
+        # Curve
+        ax1.plot(x_vals, y_vals, color="#38BDF8", lw=3.0, label=label_map, zorder=4)
+        ax1.axvline(0, color="#4B5563", lw=1.5, ls="--", zorder=1)
+        ax1.axhline(0, color="#4B5563", lw=1.5, ls="--", zorder=1)
+
+        # Min / Max assoluti (Individuazione leak instabili)
+        ax1.plot(true_min, 0, marker='v', color='#EF4444', markersize=7, zorder=5)
+        ax1.plot(true_max, 0, marker='v', color='#EF4444', markersize=7, zorder=5)
+        ax1.text(true_min, 0.05, f"Min:\n{true_min:.1f}", color='#EF4444', ha='center', va='bottom', fontsize=9,
+                 transform=ax1.get_xaxis_transform())
+        ax1.text(true_max, 0.05, f"Max:\n{true_max:.1f}", color='#EF4444', ha='center', va='bottom', fontsize=9,
+                 transform=ax1.get_xaxis_transform())
+
+        ax1.set_xlabel("Pre-activation $x = (X W_{exp})$", fontsize=10)
+        ax1.set_ylabel("Activation Output $f(x)$", color="#38BDF8", fontsize=10)
+        ax1.tick_params(axis='y', labelcolor="#38BDF8")
+
+        # PDF pre-attivazione
+        ax1_twin.hist(pre_flat, bins=150, range=(x_min, x_max), color="#F472B6", alpha=0.4, density=True, zorder=2)
+        ax1_twin.set_ylabel("Empirical PDF $p(x)$", color="#F472B6", fontsize=10)
+        ax1_twin.tick_params(axis='y', labelcolor="#F472B6")
+
+        ax1.set_title(f"L{layer_idx} - Overlay: Activation $f(x)$ & Input Distribution", color="#A78BFA", fontsize=11)
+        ax1.legend(loc="upper left", fontsize=9)
+
+        # ── Subplot 2: Post-act distribution (Sparsity check) ──
+        ax2 = axes[1]
+        post_min, post_max = np.percentile(post_flat, [0.0, 99.9])
+        ax2.hist(post_flat, bins=120, range=(post_min, max(post_max, 1e-3)), color="#F472B6", alpha=0.6, density=True)
+        ax2.axvline(0, color="#9CA3AF", lw=1, ls="--", alpha=0.5)
+        ax2.set_xlabel("Post-activation $Act(X W_{exp})$", fontsize=10)
+        ax2.set_ylabel("Density", fontsize=10)
+        ax2.set_title("Post-Activation Distribution (Sparsity check)", color="#A78BFA", fontsize=11)
+        ax2.grid(True, ls="--", alpha=0.3)
+
+        fig.tight_layout()
+        self.show_fig(fig)
+
+    # ─── TRAINING HISTORY ───────────────────────────��─────────────────────────
+
+    def plot_training_history(self):
+        """Plots losses and metrics matching the Vathos logic."""
+        losses_dict = getattr(self.model, "_losses_dict", {})
+        losses_ep = getattr(self.model, "_losses_per_epoch_dict", {})
+        metrics_ep = getattr(self.model, "_metrics_per_epoch", {})
+
+        if not losses_dict and not metrics_ep:
+            print("No training history found in this model.")
+            return
+
+        print("=== 📈 Training History ===")
+        print(
+            f"Steps: {getattr(self.model, 'steps', '-')} | Epochs: {getattr(self.model, 'epochs', '-')} | Best Loss: {getattr(self.model, 'best_loss', '-')}")
+
+        if losses_dict or losses_ep:
+            fig, ax = plt.subplots(figsize=(10, 4))
+            if losses_dict:
+                xs = list(losses_dict.keys())
+                ys = list(losses_dict.values())
+                ax.plot(xs, ys, color="#4c9be8", linewidth=0.8, alpha=0.4, label="Loss (step)")
+                # Smooth curve
+                if len(ys) > 10:
+                    win = max(2, len(ys) // 20)
+                    ys_sm = np.convolve(ys, np.ones(win) / win, mode="valid")
+                    ax.plot(xs[win - 1:], ys_sm, color="#4c9be8", linewidth=1.5, label=f"Smoothed (w={win})")
+
+            if losses_ep:
+                ex = list(losses_ep.keys())
+                ey = list(losses_ep.values())
+                ax.plot(ex, ey, color="#f5a623", linewidth=2.2, marker="o", markersize=5, label="Loss per epoch")
+
+            ax.set_title("Training Loss", color="#A78BFA")
+            ax.set_xlabel("Steps / Epochs")
+            ax.set_ylabel("Loss")
+            ax.grid(True, ls="--", alpha=0.4)
+            ax.legend(fontsize=9)
+            fig.tight_layout()
+            self.show_fig(fig)
+
+        # Plot Metrics
+        for t_idx, (m_name, m_vals) in enumerate(metrics_ep.items()):
+            fig, ax = plt.subplots(figsize=(10, 3.5))
+            col = COLORS[t_idx % len(COLORS)]
+            ax.plot(range(len(m_vals)), m_vals, color=col, linewidth=2.2, marker="o", markersize=5,
+                    label=f"{m_name} per epoch")
+            ax.set_title(f"Metric: {m_name}", color="#A78BFA")
+            ax.set_xlabel("Epochs")
+            ax.set_ylabel(m_name)
+            ax.grid(True, ls="--", alpha=0.4)
+            ax.legend(fontsize=9)
+            fig.tight_layout()
+            self.show_fig(fig)
+
+    def plot_all(self, layer_idx: int = 0, tokens: Optional[List[str]] = None):
+        """
+        Esegue un'analisi topologica e statistica completa, chiamando in sequenza
+        tutte le routine grafiche. Genera un referto end-to-end del modello.
+        """
+        print(f"\n{'=' * 70}\n🔬 INITIATING FULL TOPOLOGICAL ANALYSIS (Focus Layer {layer_idx})\n{'=' * 70}")
+        self.plot_overview()
+
+        print(f"\n[1/4] Weight Distributions (Layer {layer_idx})")
+        self.plot_weight_distributions(layer_idx=layer_idx)
+
+        print(f"\n[2/4] SVD Spectrum & Condition Numbers (Layer {layer_idx})")
+        self.plot_svd_spectrum(layer_idx=layer_idx)
+
+        if not self.captured_activations:
+            print("\n⚠️ Nessun forward pass catturato nello stato. Salto l'analisi del manifold (Attention/FFN).")
+            print("    Chiama `capture_forward(x)` prima di `plot_all()` per abilitarla.")
+        else:
+            print(f"\n[3/4] Attention Topology (Layer {layer_idx})")
+            self.plot_attention_weights(layer_idx=layer_idx, tokens=tokens)
+
+            print(f"\n[4/4] FFN Manifold & Activation Geometry f(x) (Layer {layer_idx})")
+            self.plot_ffn_manifold(layer_idx=layer_idx)
+
+        print("\n--- Training History ---")
+        self.plot_training_history()
+        print(f"\n{'=' * 70}\n✅ ANALYSIS COMPLETE\n{'=' * 70}")
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  TOP-LEVEL MACRO
+# ──────────────────────────────────────────────────────────────────────────────
+
+def study(model: nn.Module, layer_idx: int = 0, dummy_seq_len: int = 64,
+          device: Optional[Union[str, torch.device]] = None) -> VathosAnalyzer:
+    """
+    Macro globale per profilare immediatamente un modello in un Notebook o script nudo.
+    1. Assicura il corretto posizionamento e deduce il device dei tensori.
+    2. Istanzia il VathosAnalyzer.
+    3. Genera un input sintetico coerente col dominio.
+    4. Cattura il manifold (iniettando i Forward Hooks estrae i pre-logits).
+    5. Esegue tutta la suite grafica (plot_all).
+
+    Ritorna:
+        L'istanza di VathosAnalyzer configurata per ulteriori ispezioni interattive.
+    """
+    if device is None:
+        try:
+            device = next(model.parameters()).device
+        except StopIteration:
+            device = torch.device("cpu")
+            model.to(device)
+
+    print(f"🔍 Starting Rigorous Study on {type(model).__name__} (Device: {device})")
+
+    analyzer = VathosAnalyzer(model)
+
+    # Tentiamo di dedurre il vocab_size rigorosamente, con fallback
+    vocab_size = getattr(model, "vocab_size", 1000)
+    if not hasattr(model, "vocab_size") and hasattr(model, "embedder") and hasattr(model.embedder, "embedding"):
+        vocab_size = getattr(model.embedder.embedding, "num_embeddings", 1000)
+
+    # Dummy input per forzare il passaggio nei manifold locali ed estrarre la Jacobiana implicita
+    dummy_x = torch.randint(0, vocab_size, (1, dummy_seq_len), device=device)
+
+    print(f"📸 Capturing activation manifold with dummy sequence (L={dummy_seq_len})...")
+    analyzer.capture_forward(dummy_x)
+
+    # Lancia l'analisi completa concentrandosi sul layer indicato (0 di default)
+    analyzer.plot_all(layer_idx=layer_idx)
+
+    return analyzer

vathos/Transformers.py

@@ -0,0 +1,264 @@
+"""Basic Implementation of Famous Transformers Model using Vathos blocks SequenceModel.
+serves firstly as an example on how to build with Vathos SequenceModel"""
+
+from Vathos.blocks import *
+import torch.nn as nn
+from Vathos.functions import *
+
+
+class GPTTransformer(nn.Module):
+    def __init__(self, vocab_size, d_model, n_layers, n_heads, max_len, ff_expand=2, activation='gelu', dropout=0.1,
+                 pe='sinu'):
+        super(GPTTransformer, self).__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.ff_expand = ff_expand
+        self.activation = activation
+        self.dropout = dropout
+        if pe == 'sinu':
+            self.pe = True
+            self.rope = False
+        elif pe == 'rope':
+            self.pe = False
+            self.rope = True
+        elif self.pe == 'nope':
+            self.pe = False
+        else:
+            raise ValueError('pe={} not recognized'.format(pe))
+
+        self.model = SequenceModel(
+            vocab_size=vocab_size,
+            d_model=d_model,
+            n_layers=n_layers,
+            max_len=max_len,
+            pos_encoder=self.pe,
+            rope=self.rope,
+            spatial_args={"causal": True, "n_heads": n_heads, "rope": self.rope},
+            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+    def __repr__(self):
+        return self.model.__repr__()
+
+    def __str__(self):
+        return self.model.__str__()
+
+    @torch.no_grad()
+    def generate(self, *args, **kwargs):
+        return self.model.generate(*args, *kwargs)
+
+    def summary(self):
+        self.model.summary()
+
+
+class BERTTransformer(nn.Module):
+    def __init__(self, vocab_size, d_model, n_layers, n_heads, max_len, ff_expand=2, activation='gelu', dropout=0.1,
+                 pe='sinu'):
+        super(GPTTransformer, self).__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.ff_expand = ff_expand
+        self.activation = activation
+        self.dropout = dropout
+        if pe == 'sinu':
+            self.pe = True
+            self.rope = False
+        elif pe == 'rope':
+            self.pe = False
+            self.rope = True
+        elif self.pe == 'nope':
+            self.pe = False
+        else:
+            raise ValueError('pe={} not recognized'.format(pe))
+
+        self.model = SequenceModel(
+            vocab_size=vocab_size,
+            d_model=d_model,
+            n_layers=n_layers,
+            max_len=max_len,
+            pos_encoder=self.pe,
+            rope=self.rope,
+            spatial_args={"causal": False, "n_heads": n_heads, "rope": self.rope},
+            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+    def __repr__(self):
+        return self.model.__repr__()
+
+    def __str__(self):
+        return self.model.__str__()
+
+    @torch.no_grad()
+    def generate(self, *args, **kwargs):
+        raise NotImplementedError("Yet to Implement")
+
+    def summary(self):
+        self.model.summary()
+
+
+class DeiTTransformer(nn.Module):
+    def __init__(self, n_class, img_size, patch_size, d_model, n_layers, n_heads, in_chans=3, ff_expand=2,
+                 activation='gelu', dropout=0.1,
+                 pe='sinu'):
+        super().__init__()
+        self.img_size = img_size
+        self.n_class = n_class
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.ff_expand = ff_expand
+        self.activation = activation
+        self.dropout = dropout
+        if pe == 'sinu':
+            self.pe = True
+            self.rope = False
+        elif pe == 'rope':
+            self.pe = False
+            self.rope = True
+        elif self.pe == 'nope':
+            self.pe = False
+        else:
+            raise ValueError('pe={} not recognized'.format(pe))
+
+        self.model = SequenceModel(
+            vocab_size=n_class,
+            d_model=d_model,
+            n_layers=n_layers,
+            pos_encoder=self.pe,
+            rope=self.rope,
+            embedder=PatchEmbedder,
+            embedder_args={"d_model": d_model, "img_size": img_size, "patch_size": patch_size, "in_chans": in_chans},
+            unembedder=ClsHead,
+            spatial_args={"causal": False, "n_heads": n_heads, "rope": self.rope},
+            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+    def __repr__(self):
+        return self.model.__repr__()
+
+    def __str__(self):
+        return self.model.__str__()
+
+    @torch.no_grad()
+    def generate(self, *args, **kwargs):
+        raise NotImplementedError("Yet to Implement")
+
+    def summary(self):
+        self.model.summary()
+
+
+class ViTTTransformer(nn.Module):
+    def __init__(self, n_class, img_size, patch_size, d_model, n_layers, n_heads, in_chans=3, ff_expand=2,
+                 activation='gelu', dropout=0.1,
+                 pe='sinu'):
+        super().__init__()
+        self.img_size = img_size
+        self.n_class = n_class
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.ff_expand = ff_expand
+        self.activation = activation
+        self.dropout = dropout
+        if pe == 'sinu':
+            self.pe = True
+            self.rope = False
+        elif pe == 'rope':
+            self.pe = False
+            self.rope = True
+        elif self.pe == 'nope':
+            self.pe = False
+        else:
+            raise ValueError('pe={} not recognized'.format(pe))
+
+        self.model = SequenceModel(
+            vocab_size=n_class,
+            d_model=d_model,
+            n_layers=n_layers,
+            pos_encoder=self.pe,
+            rope=self.rope,
+            embedder=PatchEmbedder,
+            embedder_args={"d_model": d_model, "img_size": img_size, "patch_size": patch_size, "in_chans": in_chans},
+            unembedder=MeanClassificationHead,
+            spatial_args={"causal": False, "n_heads": n_heads, "rope": self.rope},
+            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+    def __repr__(self):
+        return self.model.__repr__()
+
+    def __str__(self):
+        return self.model.__str__()
+
+    @torch.no_grad()
+    def generate(self, *args, **kwargs):
+        raise NotImplementedError("Yet to Implement")
+
+    def summary(self):
+        self.model.summary()
+
+
+class DiTTransformer(nn.Module):
+    def __init__(self, d_model, n_layers, n_heads, ff_expand=2, activation='gelu', dropout=0.1,
+                 pe='sinu'):
+        super().__init__()
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.ff_expand = ff_expand
+        self.activation = activation
+        self.dropout = dropout
+        if pe == 'sinu':
+            self.pe = True
+            self.rope = False
+        elif pe == 'rope':
+            self.pe = False
+            self.rope = True
+        elif self.pe == 'nope':
+            self.pe = False
+        else:
+            raise ValueError('pe={} not recognized'.format(pe))
+
+        self.model = SequenceModel(
+            vocab_size=2,
+            d_model=d_model,
+            n_layers=n_layers,
+            pos_encoder=self.pe,
+            rope=self.rope,
+            embedder=Identity,
+            unembedder=Identity,
+            spatial_args={"causal": False, "n_heads": n_heads, "rope": self.rope},
+            channel_args={"expand": ff_expand, "activation": ACTIVS[activation], "depth": 2}
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+    def __repr__(self):
+        return self.model.__repr__()
+
+    def __str__(self):
+        return self.model.__str__()
+
+    @torch.no_grad()
+    def generate(self, *args, **kwargs):
+        raise NotImplementedError("Yet to Implement")
+
+    def summary(self):
+        self.model.summary()

vathos/__init__.py

@@ -0,0 +1 @@
+from Vathos.blocks import *

vathos/_analytics.py

@@ -0,0 +1,439 @@
+"""
+Vathos Analytics & Interpretability Toolkit
+===========================================
+Modulo standalone per l'analisi spettrale, la topologia dell'attenzione e
+l'interpretazione semantica dei concetti (FFN) nei modelli ModdedFormer.
+
+Progettato per Jupyter Notebooks e script Python off-Streamlit.
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+import matplotlib.pyplot as plt
+from typing import Optional, Dict, List, Tuple, Any
+import math
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  1. ANALISI MATEMATICA E SPETTRALE (Core)
+# ──────────────────────────────────────────────────────────────────────────────
+
+def compute_tensor_stats(t: torch.Tensor) -> dict:
+    """Calcola i momenti statistici e la sparsità di un tensore o manifold."""
+    f = t.detach().float().cpu()
+    v = f.numpy().flatten()
+    return {
+        "shape": tuple(t.shape),
+        "numel": t.numel(),
+        "mean": float(v.mean()),
+        "std": float(v.std()),
+        "min": float(v.min()),
+        "max": float(v.max()),
+        "l2_norm": float(np.linalg.norm(v)),
+        "sparsity": float((np.abs(v) < 1e-6).mean())
+    }
+
+
+def compute_svd_spectrum(t: torch.Tensor) -> Optional[np.ndarray]:
+    """
+    Decomposizione in Valori Singolari (SVD) per studiare il collasso dimensionale.
+    Ritorna i valori singolari sigma. Se il tensore è 3D (es. Head), media gli spettri.
+    """
+    f = t.detach().float().cpu()
+    if f.ndim == 2:
+        try:
+            return torch.linalg.svdvals(f).numpy()
+        except Exception:
+            return None
+    if f.ndim == 3:
+        results = []
+        for i in range(f.shape[0]):
+            try:
+                results.append(torch.linalg.svdvals(f[i]).numpy())
+            except Exception:
+                pass
+        return np.stack(results) if results else None
+    return None
+
+
+def compute_shannon_entropy(attn_weights: np.ndarray) -> List[float]:
+    """
+    Calcola l'Entropia di Shannon E[H] lungo l'asse delle Key per ogni Head.
+    H = - sum(p * log(p)). Valori bassi indicano forte 'certezza' (es. Induction Heads).
+    """
+    w_calc = attn_weights if attn_weights.ndim == 3 else attn_weights[np.newaxis]
+    entropies = []
+    for h in range(w_calc.shape[0]):
+        wh = w_calc[h] + 1e-12  # Epsilon per stabilità numerica nel log
+        ent = -(wh * np.log(wh)).sum(axis=-1).mean()
+        entropies.append(float(ent))
+    return entropies
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  2. ESTRAZIONE TOPOLOGICA (Hooks & Forward Passes)
+# ──────────────────────────────────────────────────────────────────────────────
+
+def get_attention_matrix(mixer: nn.Module, x: torch.Tensor) -> Optional[np.ndarray]:
+    """Ricalcola rigorosamente Softmax(QK^T / sqrt(d)) per estrarre le matrici d'attenzione."""
+    if hasattr(mixer, "get_attention_weights"):
+        with torch.no_grad():
+            return mixer.get_attention_weights(x).detach().float().cpu().numpy()
+
+    try:
+        with torch.no_grad():
+            B, L, D = x.shape
+            if hasattr(mixer, "qkv"):
+                proj = mixer.qkv
+                n_heads = mixer.n_heads
+                head_dim = mixer.head_dim
+                qkv = proj(x).view(B, L, 3, n_heads, head_dim)
+                q, k, _ = qkv.unbind(dim=2)
+                q, k = q.transpose(1, 2).float(), k.transpose(1, 2).float()
+                scale = math.sqrt(head_dim)
+                scores = torch.matmul(q, k.transpose(-2, -1)) / scale
+                if getattr(mixer, "causal", True):
+                    mask = torch.tril(torch.ones(L, L, device=x.device)).bool()
+                    scores = scores.masked_fill(~mask, float("-inf"))
+                return torch.softmax(scores, dim=-1)[0].cpu().numpy()
+            elif hasattr(mixer, "qk"):
+                proj = mixer.qk
+                n_heads = mixer.n_heads
+                head_dim = mixer.head_dim
+                qk = proj(x).view(B, L, 2, n_heads, head_dim)
+                q, k = qk.unbind(dim=2)
+                q, k = q.transpose(1, 2).float(), k.transpose(1, 2).float()
+                scale = math.sqrt(head_dim)
+                scores = torch.matmul(q, k.transpose(-2, -1)) / scale
+                if getattr(mixer, "causal", True):
+                    mask = torch.tril(torch.ones(L, L, device=x.device)).bool()
+                    scores = scores.masked_fill(~mask, float("-inf"))
+                return torch.softmax(scores, dim=-1)[0].cpu().numpy()
+    except Exception as e:
+        print(f"Errore estrazione QK: {e}")
+    return None
+
+
+def get_ffn_manifolds(cm: nn.Module, x: torch.Tensor, tau: float = None) -> Optional[Dict[str, np.ndarray]]:
+    """
+    Tratta il Channel Mixer come Memoria Associativa.
+    Ritorna Logits (XW), Attivazioni (Act(XW)) e Concept Attention.
+    """
+    if tau is None:
+        tau = math.sqrt(x.shape[-1])  # Default temperature = sqrt(d_model)
+
+    try:
+        with torch.no_grad():
+            pre_act = cm.expand(x)  # Y = XW_{expand}
+            post_act = cm.act(pre_act) if hasattr(cm, "act") else pre_act
+            concept_attn = torch.softmax(pre_act / tau, dim=-1)
+
+            return {
+                "pre_act": pre_act[0].cpu().numpy(),  # [L, M_dim]
+                "post_act": post_act[0].cpu().numpy(),  # [L, M_dim]
+                "concept_attn": concept_attn[0].cpu().numpy()  # [L, M_dim]
+            }
+    except Exception as e:
+        print(f"Errore estrazione FFN Manifold: {e}")
+    return None
+
+
+def capture_layer_manifolds(model: nn.Module, input_ids: torch.Tensor, layer_idx: int) -> Tuple[
+    torch.Tensor, torch.Tensor]:
+    """
+    Usa i forward hooks per intercettare i manifold esatti in ingresso
+    ai mixer Spaziali e di Canale per un dato strato L.
+    """
+    block = model.blocks[layer_idx]
+    sm = block.spatial_mixer
+    cm = block.channel_mixer
+
+    captured = {"spatial_x": None, "channel_x": None}
+
+    def hook_spatial(m, i):
+        captured["spatial_x"] = i[0].detach()
+
+    def hook_channel(m, i):
+        captured["channel_x"] = i[0].detach()
+
+    h1 = sm.register_forward_pre_hook(hook_spatial)
+    h2 = cm.register_forward_pre_hook(hook_channel)
+
+    try:
+        model.eval()
+        with torch.no_grad():
+            model(input_ids)
+    finally:
+        h1.remove()
+        h2.remove()
+
+    return captured["spatial_x"], captured["channel_x"]
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  3. FUNZIONI DI PLOTTING (Matplotlib Figure Generators)
+# ──────────────────────────────────────────────────────────────────────────────
+
+def set_dark_style():
+    """Opzionale: Applica il tema scuro (stile Vathos) globale a Matplotlib."""
+    plt.rcParams.update({
+        "figure.facecolor": "#0D0D1A", "axes.facecolor": "#0D0D1A", "axes.edgecolor": "#252550",
+        "axes.labelcolor": "#C0C0E0", "xtick.color": "#9CA3AF", "ytick.color": "#9CA3AF",
+        "text.color": "#D8D8F0", "grid.color": "#1E1E3A", "grid.alpha": 0.6
+    })
+
+
+def plot_weight_distributions(tensors_dict: Dict[str, torch.Tensor], bins: int = 80) -> Tuple[plt.Figure, np.ndarray]:
+    n = len(tensors_dict)
+    cols = min(n, 3)
+    rows = math.ceil(n / cols)
+    fig, axes = plt.subplots(rows, cols, figsize=(5 * cols, 3.5 * rows))
+    if n == 1:
+        axes = np.array([[axes]])
+    elif rows == 1:
+        axes = axes.reshape(1, -1)
+
+    colors = ["#A78BFA", "#60A5FA", "#34D399", "#F472B6", "#FB923C", "#FBBF24"]
+
+    for idx, (name, t) in enumerate(tensors_dict.items()):
+        ax = axes[idx // cols][idx % cols]
+        v = t.detach().float().cpu().numpy().flatten()
+        stats = compute_tensor_stats(t)
+        ax.hist(v, bins=bins, color=colors[idx % len(colors)], alpha=0.8, linewidth=0)
+        ax.set_title(f"{name}\n$\mu$={stats['mean']:.3e}  $\sigma$={stats['std']:.3e}", fontsize=9)
+        ax.grid(True, ls="--", alpha=0.4)
+
+    for idx in range(n, rows * cols): axes[idx // cols][idx % cols].set_visible(False)
+    fig.suptitle("Distribuzioni di Densità", fontsize=12)
+    fig.tight_layout()
+    return fig, axes
+
+
+def plot_svd_spectra(tensors_dict: Dict[str, torch.Tensor], log_scale: bool = True) -> Tuple[plt.Figure, np.ndarray]:
+    eligible = {k: v for k, v in tensors_dict.items() if v.ndim >= 2}
+    n = len(eligible)
+    if n == 0: raise ValueError("Nessun tensore 2D+ fornito per la SVD.")
+
+    cols = min(n, 3)
+    rows = math.ceil(n / cols)
+    fig, axes = plt.subplots(rows, cols, figsize=(5 * cols, 3.5 * rows))
+    if n == 1:
+        axes = np.array([[axes]])
+    elif rows == 1:
+        axes = axes.reshape(1, -1)
+
+    colors = ["#A78BFA", "#60A5FA", "#34D399", "#F472B6", "#FB923C", "#FBBF24"]
+
+    for idx, (name, t) in enumerate(eligible.items()):
+        ax = axes[idx // cols][idx % cols]
+        sv = compute_svd_spectrum(t)
+        if sv is None: continue
+
+        color = colors[idx % len(colors)]
+        if sv.ndim == 2:
+            for h in range(sv.shape[0]):
+                ax.plot(sv[h], color=color, alpha=0.2, linewidth=0.8)
+            ax.plot(sv.mean(0), color=color, linewidth=2, label="Mean $\Sigma$")
+            ax.legend()
+        else:
+            ax.plot(sv, color=color, linewidth=1.5, marker=".", markersize=3)
+
+        cond = sv.flatten()[0] / (sv.flatten()[-1] + 1e-12)
+        r_eff = int((sv.flatten() > sv.flatten()[0] * 1e-3).sum())
+        ax.set_title(f"{name}\n$\kappa={cond:.1f}$ | Rango Effettivo $\\approx {r_eff}$", fontsize=9)
+        if log_scale: ax.set_yscale("log")
+        ax.grid(True, ls="--", which="both", alpha=0.4)
+
+    for idx in range(n, rows * cols): axes[idx // cols][idx % cols].set_visible(False)
+    fig.tight_layout()
+    return fig, axes
+
+
+def plot_attention_topology(attn_weights: np.ndarray, tokens: Optional[List[str]] = None) -> Tuple[
+    plt.Figure, np.ndarray]:
+    """Plotta la matrice $L \times L$ per le varie Head d'attenzione."""
+    if attn_weights.ndim == 2: attn_weights = attn_weights[np.newaxis]
+    H, L_q, L_k = attn_weights.shape
+    cols = min(H, 4)
+    rows = math.ceil(H / cols)
+
+    cell_size = max(0.2, min(0.6, 12.0 / L_q)) if tokens else 0.5
+    fig_w = max(4 * cols, cell_size * L_q * cols)
+    fig_h = max(3.5 * rows, cell_size * L_q * rows + 1)
+
+    fig, axes = plt.subplots(rows, cols, figsize=(fig_w, fig_h))
+    axes = np.array(axes).flatten() if H > 1 else [axes]
+
+    clean_tokens = [t.replace('Ġ', '').replace(' ', '') for t in tokens] if tokens else None
+
+    for h in range(H):
+        ax = axes[h]
+        w = attn_weights[h]
+        im = ax.imshow(w, aspect="auto", cmap="magma", vmin=0, vmax=w.max() + 1e-9)
+        ax.set_title(f"Head {h}", fontsize=10)
+
+        if clean_tokens and len(clean_tokens) == L_q:
+            ax.set_xticks(range(L_k))
+            ax.set_yticks(range(L_q))
+            fs = max(5, min(10, int(300 / L_q)))
+            ax.set_xticklabels(clean_tokens, rotation=90, fontsize=fs)
+            ax.set_yticklabels(clean_tokens, fontsize=fs)
+        else:
+            ax.set_xlabel("Key Position (Source)")
+            ax.set_ylabel("Query Position (Target)")
+
+        fig.colorbar(im, ax=ax, shrink=0.8)
+
+    for h in range(H, len(axes)): axes[h].set_visible(False)
+    fig.tight_layout()
+    return fig, axes
+
+
+def plot_ffn_concept_manifold(matrix: np.ndarray, title: str, tokens: Optional[List[str]] = None,
+                              max_neurons: int = 128) -> Tuple[plt.Figure, plt.Axes]:
+    """Plotta la mappa $L \times M$ di attivazione dei concetti per il Channel Mixer."""
+    mat_to_plot = matrix[:, :max_neurons]
+    rows_c, cols_c = mat_to_plot.shape
+
+    fig_w = max(8, cols_c * 0.15) if tokens else 10
+    fig_h = max(6, rows_c * 0.3) if tokens else 6
+
+    fig, ax = plt.subplots(figsize=(min(fig_w, 20), min(fig_h, 15)))
+    im = ax.imshow(mat_to_plot, aspect="auto", cmap="magma")
+
+    ax.set_title(f"{title} (Primi {cols_c} neuroni su {matrix.shape[1]})", pad=15)
+    ax.set_ylabel("Tokens (Queries)")
+    ax.set_xlabel("FFN Neurons (Latent Concepts)")
+
+    if tokens and len(tokens) == rows_c:
+        clean_tokens = [t.replace('Ġ', '').replace(' ', '') for t in tokens]
+        fs = max(5, min(10, int(400 / rows_c)))
+        ax.set_yticks(range(rows_c))
+        ax.set_yticklabels(clean_tokens, fontsize=fs)
+
+    fig.colorbar(im, ax=ax)
+    fig.tight_layout()
+    return fig, ax
+
+
+def plot_architecture_scalars(model: nn.Module) -> Tuple[Optional[plt.Figure], Optional[plt.Figure]]:
+    """Genera i plot per gli scalari strutturali: ZeroSkip e Skip-connection Lambdas."""
+    fig_skip, fig_zero = None, None
+
+    # 1. Skip Lambdas
+    if hasattr(model, "skip_lambdas") and model.skip_lambdas:
+        lambdas = {k: float(v.detach()) for k, v in model.skip_lambdas.items()}
+        fig_skip, ax = plt.subplots(figsize=(max(6, len(lambdas) * 1.4), 4))
+        keys, vals = list(lambdas.keys()), list(lambdas.values())
+        ax.bar(range(len(keys)), vals, alpha=0.85, color="#60A5FA")
+        ax.set_xticks(range(len(keys)))
+        ax.set_xticklabels([k.replace("route_", "").replace("_to_", "→") for k in keys], rotation=30, ha="right")
+        ax.axhline(0, color="gray", linewidth=0.8, ls="--")
+        ax.set_title("Skip-connection Routing Gates ($\lambda$)")
+        fig_skip.tight_layout()
+
+    # 2. ZeroSkip Parameters
+    if hasattr(model, "zeroskip_params") and model.zeroskip:
+        vals = [float(p.detach()) for p in model.zeroskip_params]
+        fig_zero, ax = plt.subplots(figsize=(max(6, len(vals) * 0.7), 4))
+        ax.plot(vals, marker="o", color="#FBBF24", linewidth=1.5)
+        ax.axhline(0, color="gray", linewidth=0.8, ls="--")
+        ax.set_xlabel("Layer Index")
+        ax.set_ylabel("ZeroSkip $\\alpha$")
+        ax.set_title("ZeroSkip Coefficients per Layer")
+        ax.grid(True, ls="--", alpha=0.4)
+        fig_zero.tight_layout()
+
+    return fig_skip, fig_zero
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+#  4. HIGH-LEVEL API PER JUPYTER (Sintesi per Speedrun & ARC)
+# ──────────────────────────────────────────────────────────────────────────────
+
+def analyze_layer_semantics(model: nn.Module, tokenizer: Any, prompt: str, layer_idx: int,
+                            add_bos: bool = True, ffn_tau: float = None,
+                            show_plots: bool = True) -> Dict[str, Any]:
+    """
+    Funzione master per i Notebook. Esegue l'analisi completa del layer indicato
+    su un prompt specifico, restituendo statistiche e plottando i grafici.
+
+    Uso:
+        results = analyze_layer_semantics(model, tokenizer, "Test sequence", layer_idx=4)
+        print("Sparsità FFN:", results["ffn_sparsity"])
+    """
+    device = next(model.parameters()).device
+    full_prompt = (tokenizer.bos_token if add_bos and tokenizer.bos_token else "") + prompt
+    inputs = tokenizer(full_prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"].to(device)
+    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
+
+    # 1. Forward Pass con intercettazione Manifold
+    sm_x, cm_x = capture_layer_manifolds(model, input_ids, layer_idx)
+    if sm_x is None:
+        raise RuntimeError(f"Hook fallito. Il layer {layer_idx} non è stato eseguito.")
+
+    block = model.blocks[layer_idx]
+    sm = block.spatial_mixer
+    cm = block.channel_mixer
+
+    d_model = model.d_models[layer_idx] if hasattr(model, "d_models") else sm_x.shape[-1]
+    if ffn_tau is None: ffn_tau = math.sqrt(d_model)
+
+    # 2. Estrazione Matrici
+    attn_w = get_attention_matrix(sm, sm_x)
+    ffn_manifolds = get_ffn_manifolds(cm, cm_x, tau=ffn_tau)
+
+    results = {
+        "tokens": tokens,
+        "X_spatial": sm_x,
+        "X_channel": cm_x,
+        "attention_matrix": attn_w,
+        "ffn_manifolds": ffn_manifolds,
+    }
+
+    # 3. Calcoli Spettrali e Statistiche (Information Theory)
+    if attn_w is not None:
+        results["attention_entropy"] = compute_shannon_entropy(attn_w)
+
+    if ffn_manifolds is not None:
+        Y_act = ffn_manifolds["post_act"]
+        stats = compute_tensor_stats(torch.tensor(Y_act))
+        results["ffn_sparsity"] = stats["sparsity"]
+        results["ffn_l2_norm"] = stats["l2_norm"]
+
+        try:
+            svd_vals = np.linalg.svd(Y_act, compute_uv=False)
+            results["ffn_eff_rank"] = int(np.sum(svd_vals > svd_vals[0] * 1e-3))
+            results["ffn_svd_spectrum"] = svd_vals
+        except np.linalg.LinAlgError:
+            pass
+
+    # 4. Rendering Opzionale
+    if show_plots:
+        print(f"--- ANALISI DEL MANIFOLD | LAYER {layer_idx} ---")
+
+        if attn_w is not None:
+            print("\n>> Topologia Spaziale dell'Attenzione (QK)")
+            fig_attn, _ = plot_attention_topology(attn_w, tokens)
+            plt.show()
+            print(f"   E[H] (Entropia per Head): {[f'{e:.3f}' for e in results['attention_entropy']]}")
+
+        if ffn_manifolds is not None:
+            print("\n>> Channel Expansion (Memory Concepts)")
+            print(f"   Sparsità (Soglia < 1e-6): {results.get('ffn_sparsity', 0):.2%}")
+            print(f"   Rango Effettivo (SVD): {results.get('ffn_eff_rank', 'N/A')} su {Y_act.shape[1]} M_dim")
+
+            fig_ffn, _ = plot_ffn_concept_manifold(ffn_manifolds["post_act"], "FFN $Act(X W_{expand})$", tokens)
+            plt.show()
+
+            if "ffn_svd_spectrum" in results:
+                fig_svd, ax_svd = plt.subplots(figsize=(6, 3))
+                ax_svd.plot(results["ffn_svd_spectrum"], color="red", marker=".")
+                ax_svd.set_yscale("log")
+                ax_svd.set_title("Spettro $\Sigma$ del Manifold dei Concetti (FFN)")
+                ax_svd.grid(True, ls="--", alpha=0.5)
+                plt.show()
+
+    return results

vathos/_basics.py

@@ -0,0 +1,1134 @@
+"""
+This modules serves as base for the Torch Basic model Structure, and profilers.
+Class Layer is the base of all Vathos layers, and ereditate from torch.nn.Module class, adding only the structures that help
+profiling, visualization, and debugging.
+Everything built with Layer is totally compatible with torch native modules.
+"""
+
+import numpy as np
+import torch.nn as nn
+from typing import Callable
+from Vathos.functions import *
+from timeit import default_timer as timer
+from collections import OrderedDict, defaultdict
+import re
+import math
+from torch import Tensor
+
+ACTIVS = {
+    'tanh': nn.Tanh,
+    'sigmoid': nn.Sigmoid,
+    'relu': nn.ReLU,
+    'gelu': nn.GELU,
+    'elu': nn.ELU,
+    'lrelu': nn.LeakyReLU,
+    'leaky_relu': nn.LeakyReLU,
+}
+
+
+class VathosConfig:
+    """Global configuration state."""
+    _COMPILABLE = False  # True = production, torch.compile safe
+    _PROFILE_BATCHED = True  # True = times divided by batch size
+    _GLOBAL_PROFILE = True  # True = always profile in debug mode
+
+
+def set_vathos_mode(mode: str):
+    """
+    Switch Vathos mode globally.
+
+    'debug'      — profiling enabled, torch.compile unfriendly
+    'production' — profiling disabled, torch.compile safe
+    """
+    if mode.lower() == "production":
+        VathosConfig._COMPILABLE = True
+        print(f"{SEC}Vathos:{RES} Switched to {GOOD}PRODUCTION{RES} mode.")
+    elif mode.lower() == "debug":
+        VathosConfig._COMPILABLE = False
+        print(f"{SEC}Vathos:{RES} Switched to {NUM}DEBUG{RES} mode.")
+    else:
+        raise ValueError("Mode must be 'production' or 'debug'")
+
+
+# =============================================================================
+# BASE LAYER
+# =============================================================================
+
+class Layer(nn.Module):
+    """
+    ─────────────────────────────────────────────────────────────
+    class MyBlock(Layer):
+        def __init__(self, ...):
+            super().__init__()
+            # define sub-modules here
+            self.lin = nn.Linear(...)
+            self.attn = MultiheadAttentionMixer(...)
+
+        def forward(self, x):
+z            return self.lin(x)
+    ─────────────────────────────────────────────────────────────
+
+    DO NOT override __call__. Profiling is handled via native forward hooks,
+    which are registered automatically in debug mode and completely absent
+    in production mode — zero overhead, full torch.compile compatibility.
+
+    IMPORTANT: gli hook sono registrati al __init__. Se il modello è costruito
+    PRIMA di chiamare set_vathos_mode('production'), gli hook sopravvivono.
+    Soluzioni:
+      1) Chiama set_vathos_mode('production') PRIMA di istanziare moduli.
+      2) Oppure chiama model.strip_profiling_hooks() dopo lo switch.
+    """
+
+    _is_vathos_layer = True
+
+    def __init__(self):
+        super().__init__()
+        self.complexity = "O(1)"
+        self.__name__ = self.__class__.__name__
+
+        self._timer_unbatched = not VathosConfig._PROFILE_BATCHED
+        self._tstart = 0.0
+        self._tend = 0.0
+        self._time = 0.0
+        self._times = []
+
+        self._sublayers = None
+
+        # Hooks are registered ONCE at construction in debug mode.
+        # In production mode, no hooks exist — no overhead whatsoever.
+        if not VathosConfig._COMPILABLE and VathosConfig._GLOBAL_PROFILE:
+            self._register_profiling_hooks()
+
+    # -------------------------------------------------------------------------
+    # Profiling hooks — native nn.Module mechanism, __call__ is never touched
+    # -------------------------------------------------------------------------
+
+    def _register_profiling_hooks(self):
+        self.register_forward_pre_hook(self._pre_hook)
+        self.register_forward_hook(self._post_hook)
+
+    def _pre_hook(self, module, args):
+        self._tstart = timer()
+
+    def _post_hook(self, module, args, output):
+        self._tend = timer()
+        bs = 1
+        if args and isinstance(args[0], torch.Tensor):
+            bs = args[0].shape[0]
+        div = bs if not self._timer_unbatched else 1
+        self._time = (self._tend - self._tstart) / div
+        self._times.append(self._time)
+
+    # -------------------------------------------------------------------------
+    # Sublayer registry — used by profile()
+    # -------------------------------------------------------------------------
+
+    def register_sublayers(self):
+        if self._sublayers is None:
+            self._sublayers = dict()
+
+        def get_unique_name(base_name, existing_names):
+            if base_name not in existing_names:
+                return base_name
+            counter = 1
+            while f"{base_name}_{counter}" in existing_names:
+                counter += 1
+            return f"{base_name}_{counter}"
+
+        def collect_layers(module, level=0):
+            layers = []
+            for name, child in module.named_children():
+                if getattr(child, "_is_vathos_layer", False):
+                    layers.append((name, child, level))
+                    layers.extend(collect_layers(child, level=level + 1))
+                elif isinstance(child, nn.ModuleList):
+                    for i, item in enumerate(child):
+                        if getattr(item, "_is_vathos_layer", False):
+                            layers.append((f"{name}[{i}]", item, level))
+                            layers.extend(collect_layers(item, level=level + 1))
+                elif isinstance(child, nn.Module):
+                    layers.extend(collect_layers(child, level=level))
+            return layers
+
+        for original_name, layer, level in collect_layers(self):
+            class_name = getattr(layer, "__name__", type(layer).__name__)
+            unique_name = get_unique_name(class_name, self._sublayers.keys())
+            self._sublayers[unique_name] = {"layer": layer, "level": level}
+
+    # -------------------------------------------------------------------------
+    # Utilities
+    # -------------------------------------------------------------------------
+
+    def get_mean_execution_time(self) -> float:
+        return float(np.mean(self._times)) if self._times else 0.0
+
+    def strip_profiling_hooks(self):
+        """Rimuove gli hook profile da questo modulo e da tutti i sotto-Layer.
+
+        Utile se il modello è stato istanziato in debug mode e poi si vuole
+        passare a production senza ricostruirlo. Idempotente.
+        """
+        for module in self.modules():
+            if isinstance(module, Layer):
+                module._forward_pre_hooks.clear()
+                module._forward_hooks.clear()
+        return self
+
+    def has_custom_generate(self) -> bool:
+        return type(self).generate is not Layer.generate
+
+    def generate(self, *args, **kwargs):
+        """Override in subclasses that support autoregressive generation."""
+        return None
+
+    def clear_times(self):
+        """Reset profiling history."""
+        self._times.clear()
+
+    def profile(self, maxlevel=100, avg=False, plot=False, plot_level=1):
+        if VathosConfig._COMPILABLE:
+            print(f"{BAD}Cannot profile in PRODUCTION mode.{RES} Run set_vathos_mode('debug') first.")
+            return
+
+        batched = not self._timer_unbatched
+        print(
+            f"Layer {NUM}{self.__name__}{RES} Times Profile (batched: {GOOD if batched else BAD}{batched}{RES}) (averaged: {GOOD if avg else BAD}{avg}{RES}):")
+
+        grouped_layers = OrderedDict()
+        order_map = {}
+        order_counter = 0
+
+        if self._sublayers:
+            for sublayer_name, sublayer_info in self._sublayers.items():
+                layer = sublayer_info['layer']
+                level = sublayer_info['level']
+                if level >= maxlevel: continue
+
+                match = re.match(r'^(.+?)_(\d+)$', sublayer_name)
+                base_name = match.group(1) if match else sublayer_name
+                key = (base_name, level)
+
+                if key not in order_map:
+                    order_map[key] = order_counter
+                    order_counter += 1
+                    grouped_layers[key] = []
+                grouped_layers[key].append((sublayer_name, layer))
+
+        if not avg:
+            if self._sublayers:
+                for sublayer_name, sublayer_info in self._sublayers.items():
+                    if sublayer_info['level'] < maxlevel:
+                        indent = "    " * sublayer_info['level']
+                        t = sublayer_info['layer'].get_mean_execution_time()
+                        print(f"{indent}- {NUM}{sublayer_name}{RES}: {t * 1000:.2f}ms")
+        else:
+            for (base_name, level), layers_list in grouped_layers.items():
+                indent = "    " * level
+                times = [l.get_mean_execution_time() for _, l in layers_list if len(l._times) > 0]
+                if times:
+                    avg_time = np.mean(times)
+                    print(
+                        f"{indent}- {NUM}{base_name}_avg{RES}: {avg_time * 1000:.2f}ms {SEC}(x{len(layers_list)}){RES}")
+                else:
+                    print(f"{indent}- {NUM}{base_name}_avg{RES}: no time recorded")
+
+        if plot:
+            try:
+                import matplotlib.pyplot as plt
+                level_layers = defaultdict(list)
+
+                for (base_name, level), layers_list in grouped_layers.items():
+                    if level == plot_level:
+                        for _, layer in layers_list:
+                            if len(layer._times) > 0:
+                                level_layers[base_name].append(np.mean(layer._times))
+
+                    elif level < plot_level:
+
+                        pass
+
+                if level_layers:
+                    layer_times = {k: sum(v) for k, v in level_layers.items()}
+                    sorted_layers = sorted(layer_times.items(), key=lambda x: x[1], reverse=True)
+                    labels = [x[0] for x in sorted_layers]
+                    times = [x[1] * 1000 for x in sorted_layers]
+
+                    fig, ax = plt.subplots(figsize=(12, 8))
+                    wedges, texts, autotexts = ax.pie(times, autopct='%1.1f%%', startangle=90)
+                    ax.legend(wedges, [f'{l}: {t:.2f}ms' for l, t in zip(labels, times)],
+                              loc="center left", bbox_to_anchor=(1, 0, 0.5, 1))
+                    ax.set_title(f'{self.__name__} - Level {plot_level}')
+                    plt.tight_layout()
+                    plt.show()
+                else:
+                    print(f"{SEC}No data for plot level {plot_level}{RES}")
+            except ImportError:
+                print(f"{BAD}Matplotlib missing{RES}")
+
+    def __repr__(self):
+        return f"{SEC}Vathos{RES}: " + super().__repr__()
+
+
+class Builder:
+    def __init__(self, layer, **params):
+        self.layer = layer
+        self.params = params
+        for att in layer.__dict__:
+            setattr(self, att, getattr(layer, att))
+
+    def __call__(self, *args):
+        return self.layer(*args, **self.params)
+
+
+class tWrapper(Layer):
+    __name__ = "tWrapper"
+
+    def __init__(self, module: nn.Module):
+        super(tWrapper, self).__init__()
+        self.module = module
+
+    def forward(self, *args, **kwargs):
+        self.module(*args, **kwargs)
+
+
+class Identity(Layer):
+    __name__ = "Identity"
+    __complexity__ = "O(1)"
+
+    def __init__(self, *args, **kwargs):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+
+
+class Skip(Layer):
+    __name__ = "Skip"
+    __complexity__ = "O(1)"
+
+    def __init__(self, layer):
+        super(Skip, self).__init__()
+        self.layer = layer
+
+    def forward(self, x):
+        return self.layer(x) + x
+
+
+class LearntSkip(Layer):
+    __name__ = "Skip"
+    __complexity__ = "O(1)"
+
+    def __init__(self, layer):
+        super().__init__()
+        self.layer = layer
+        self.w = nn.Parameter(torch.tensor([1.0, 1.0]), requires_grad=True)
+
+    def forward(self, x):
+        return self.layer(x) * self.w[0] + x * self.w[1]
+
+
+class IdentityMixer(Layer):
+    __name__ = "Identity"
+    __complexity__ = "O(1)"
+
+    def __init__(self, d_model):
+        super(IdentityMixer, self).__init__()
+        self.d_model = d_model
+
+    def forward(self, x):
+        return x
+
+
+class LPadder(Layer):
+    __name__ = "LPadder"
+    __complexity__ = "O(k d)"
+
+    def __init__(self, right=0, left=0, element=0):
+        super(LPadder, self).__init__()
+        self.right = right
+        self.left = left
+        self.element = element
+
+    def forward(self, x):
+        return F.pad(x, (0, 0, self.left, self.right), mode="constant", value=self.element)
+
+
+class dPadder(Layer):
+    __name__ = "LPadder"
+    __complexity__ = "O(k L)"
+
+    def __init__(self, up, down=0, element=0):
+        super(dPadder, self).__init__()
+        self.right = up
+        self.left = down
+        self.element = element
+
+    def forward(self, x):
+        return F.pad(x, (self.left, self.right), mode="constant", value=self.element)
+
+
+class LUnPadder(Layer):
+    __name__ = "LUnPadder"
+    __complexity__ = "O(k d)"
+
+    def __init__(self, right=0, left=0):
+        super().__init__()
+        self.right = right
+        self.left = left
+
+    def forward(self, x):
+        return x[:, self.left:-self.right, :]
+
+
+class dUnPadder(Layer):
+    __name__ = "LUnPadder"
+    __complexity__ = "O(k d)"
+
+    def __init__(self, right=0, left=0):
+        super().__init__()
+        self.right = right
+        self.left = left
+
+    def forward(self, x):
+        return x[:, :, self.left:-self.right]
+
+
+class Linear(Layer):
+    __name__ = 'Linear'
+
+    def __init__(self, input_dim, output_dim, bias=True, **kwargs):
+        super().__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.bias = bias
+        self.linear = nn.Linear(input_dim, output_dim, bias=bias)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class BottleneckRepeatedLinear(nn.Module):
+    def __init__(self, in_features, out_features, num_repeats, bias=False, scale_variance=False):
+        super().__init__()
+
+        assert out_features % num_repeats == 0, "out_features must be strictly divisible by num_repeats"
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.num_repeats = num_repeats
+        self.bottleneck_dim = out_features // num_repeats
+        self.scale_variance = scale_variance
+
+        self.weight = nn.Parameter(torch.Tensor(self.bottleneck_dim, in_features))
+
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(self.bottleneck_dim))
+        else:
+            self.register_parameter('bias', None)
+
+        self.init_parameters()
+
+    def init_parameters(self):
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, x):
+        y = F.linear(x, self.weight, self.bias)
+
+        repeat_shape = [1] * (y.dim() - 1) + [self.num_repeats]
+
+        out = y.repeat(*repeat_shape)
+
+        if self.scale_variance:
+            out = out / math.sqrt(self.num_repeats)
+
+        return out
+
+
+class ProductLinear(Layer):
+    __name__ = 'Linear'
+
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.weight1 = nn.Parameter(0.1 * torch.randn(output_dim, input_dim))
+        self.weight2 = nn.Parameter(torch.ones(output_dim, input_dim + 0.01 * torch.randn(output_dim, input_dim)))
+
+    def forward(self, x):
+        return F.linear(x, self.weight1 * self.weight2)
+
+
+class DoubleLinear(nn.Module):
+    def __init__(self, in_features: int, out_features: int,
+                 device=None, dtype=None):
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+
+        self.weight1 = nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs))
+        nn.init.xavier_uniform_(self.weight1)
+
+        self.weight2 = nn.Parameter(torch.empty((in_features, in_features), **factory_kwargs))
+        nn.init.xavier_uniform_(self.weight2)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        W_eff = self.weight1 @ self.weight2
+        return F.linear(x, W_eff)
+
+
+class LowRankLinear(Layer):
+    __name__ = 'Linear'
+
+    def __init__(self, input_dim, output_dim, rank=16, bias=False):
+        super().__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.bias = bias
+        self.l1 = nn.Linear(input_dim, rank, bias=bias)
+        self.l2 = nn.Linear(rank, output_dim, bias=bias)
+
+    def forward(self, x):
+        return self.l2(self.l1(x))
+
+
+class UnbiasedLinear(Layer):
+    __name__ = "UnbiasedLinear"
+    __complexity__ = "O(L d^2)"
+
+    def __init__(self, input_features, output_features):
+        super(UnbiasedLinear, self).__init__()
+        self.linear = nn.Linear(input_features, output_features, bias=False)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class SwiGLU(Layer):
+    gated = True
+    __name__ = "SwiGLU"
+    __complexity__ = "O(L)"
+
+    def forward(self, x: torch.Tensor):
+        x, gate = x.chunk(2, dim=-1)
+        return x * F.silu(gate)
+
+
+class ReLU2(Layer):
+    gated = False
+    __name__ = "ReLU^2"
+    __complexity__ = "O(L)"
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_pos = F.relu(x)
+        return x_pos * x_pos
+
+
+class LeakyReLU2(Layer):
+    gated = False
+    __name__ = "LeakyReLU^2"
+    __complexity__ = "O(L)"
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_pos = F.leaky_relu(x, 0.5)
+        return x_pos * x_pos
+
+
+class PSiLU2(Layer):
+    gated = False
+    __name__ = "SiLU^2"
+    __complexity__ = "O(L)"
+
+    def __init__(self):
+        super().__init__()
+        self.param = nn.Parameter(torch.tensor([5.0, 2.0]), requires_grad=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_pos = self.param[0] * F.silu(x / self.param[1])
+        return x_pos * x_pos
+
+
+class LeLU2(Layer):
+    gated = False
+    __name__ = "ReLU^2"
+    __complexity__ = "O(L)"
+
+    def __init__(self):
+        super().__init__()
+        self.a = nn.Parameter(torch.randn(1) - 0.05)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_pos = F.relu(x - self.a)
+        return x_pos * x_pos
+
+
+class PReLU2(Layer):
+    def __init__(self, num_parameters=1):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(num_parameters))
+
+    def forward(self, x: torch.Tensor):
+        return self.alpha * torch.square(F.relu(x))
+
+
+class UDLPReLU2(Layer):
+    """Unbiased Dual Layer Perceptron, Relu^2 Activation"""
+
+    def __init__(self, d_model, expand, dropout=0.075):
+        super().__init__()
+        self.expand = nn.Linear(d_model, d_model * expand, bias=False)
+        self.contract = nn.Linear(d_model * expand, d_model, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+        nn.init.kaiming_normal_(self.expand.weight, mode='fan_in', nonlinearity='relu')
+        nn.init.zeros_(self.contract.weight)
+
+    def forward(self, x):
+        return self.dropout(self.contract(torch.relu(self.expand(x)).square()))
+
+
+class MLP(Layer):
+    __name__ = "MLP"
+    __complexity__ = "O(depth L d^2)"
+
+    def __init__(self, d_model: int, depth: int, expand: int, activation: Callable, dropout=0.1):
+        super().__init__()
+        hidden_dim = d_model * expand
+        self.d_model = d_model
+        self.depth = depth
+        self.expand = expand
+        self.activation = activation
+        self.dropout = nn.Dropout(dropout)
+        layers = []
+        for i in range(depth):
+            if i == 0:
+                in_dim = d_model
+                out_dim = hidden_dim
+            elif i == depth - 1:
+                in_dim = hidden_dim
+                out_dim = d_model
+            else:
+                in_dim = hidden_dim
+                out_dim = hidden_dim
+
+            if hasattr(activation, 'gated') and i < depth - 1:
+                out_dim = out_dim * 2
+
+            layers.append(nn.Linear(in_dim, out_dim, bias=True))
+            if i < depth - 1:
+                layers.append(activation())
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor):
+        return self.dropout(self.layers(x))
+
+
+class DLPGelu(Layer):
+    def __init__(self, d_model, expand, dropout=0.1):
+        super().__init__()
+        self.expand = nn.Linear(d_model, d_model * expand, bias=True)
+        self.contract = nn.Linear(d_model * expand, d_model, bias=True)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.dropout(self.contract(F.gelu(self.expand(x))))
+
+
+class DLPSoftmax(Layer):
+    def __init__(self, d_model, m, dropout=0.1, copy=False):
+        super().__init__()
+        self.expand = nn.Linear(d_model, m, bias=False)
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.contract = nn.Linear(m, d_model, bias=False)
+
+        self.scaler = nn.Parameter(torch.tensor([1 / math.sqrt(d_model)]))
+
+        self.dropout = nn.Dropout(dropout)
+
+        if copy:
+            with torch.no_grad():
+                self.contract.weight.copy_(self.expand.weight.t())
+
+    def forward(self, x):
+        q = self.q_proj(x)
+        logits = self.expand(q * self.scaler)
+
+        attn_weights = logits.softmax(dim=-1)
+        attn_weights = self.dropout(attn_weights)
+
+        return self.contract(attn_weights)
+
+
+class FlashSDLP(Layer):
+    def __init__(self, d_model, m, num_heads, dropout=0.1, outproj=False):
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.m = m
+        self.dropout_p = dropout
+
+        self.M1 = nn.Parameter(torch.randn(num_heads, m, self.head_dim))
+        self.M2 = nn.Parameter(torch.randn(num_heads, m, self.head_dim))
+
+        self.scaler = nn.Parameter(torch.tensor([1.0]))
+
+        if outproj:
+            self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        else:
+            self.out_proj = nn.Identity()
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_normal_(self.M1)
+        nn.init.xavier_normal_(self.M2)
+
+        if isinstance(self.out_proj, nn.Linear):
+            nn.init.xavier_uniform_(self.out_proj.weight)
+
+    def forward(self, x):
+        batch_size, seq_len, _ = x.shape
+
+        q = x.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        k = self.M1.unsqueeze(0).expand(batch_size, -1, -1, -1)
+        v = self.M2.unsqueeze(0).expand(batch_size, -1, -1, -1)
+
+        attn_output = F.scaled_dot_product_attention(
+            query=q * self.scaler,
+            key=k,
+            value=v,
+            dropout_p=self.dropout_p if self.training else 0.0,
+            is_causal=False
+        )
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
+
+        return self.out_proj(attn_output)
+
+
+class DLPSwiGLU(Layer):
+    def __init__(self, d_model, expand, dropout=0.00):
+        super().__init__()
+        self.expand = nn.Linear(d_model, d_model * expand * 2, bias=True)
+        self.contract = nn.Linear(d_model * expand, d_model, bias=True)
+        self.activation = SwiGLU()
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.dropout(self.contract(self.activation(self.expand(x))))
+
+
+class LowRankGatedDLP(Layer):
+    def __init__(self, d_model, expand, dropout=0.00, activation=nn.SiLU, rank=64, gate_activation=None):
+        super().__init__()
+        self.ga = gate_activation if gate_activation is not None else Identity()
+        self.act = activation()
+        self.M = d_model * expand
+        self.rank = rank
+        self.expand = nn.Linear(d_model, d_model * expand + rank, bias=False)
+        self.contract = nn.Linear(d_model * expand, d_model, bias=False)
+        self.rexp = nn.Linear(rank, d_model * expand, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        out1, gate1 = self.expand(x).split([self.M, self.rank], dim=-1)
+        return self.contract(self.act(out1 * self.ga(self.rexp(gate1))))
+
+
+class UDLPSwiGLU(Layer):
+    def __init__(self, d_model, expand, dropout=0.00):
+        super().__init__()
+        self.expand = nn.Linear(d_model, d_model * expand * 2, bias=False)
+        self.contract = nn.Linear(d_model * expand, d_model, bias=False)
+        self.activation = SwiGLU()
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.dropout(self.contract(self.activation(self.expand(x))))
+
+
+class VariableUDLP(Layer):
+    def __init__(self, d_model, d_output, M, activation=ReLU2, dropout=0.00):
+        super().__init__()
+        self.expand = nn.Linear(d_model, M, bias=False)
+        self.contract = nn.Linear(M, d_output, bias=False)
+        self.activation = activation()
+        self.dropout = nn.Dropout(dropout)
+
+    def _init_weights(self):
+        # Identity-at-init: expand orthogonal, contract → 0 ⇒ FFN contributo nullo.
+        torch.nn.init.orthogonal_(self.expand.weight)
+        torch.nn.init.zeros_(self.contract.weight)
+
+    def forward(self, x):
+        return self.dropout(self.contract(self.activation(self.expand(x))))
+
+
+class VariableUDLP_Attention(Layer):
+    def __init__(self, d_model, d_output, M, activation=ReLU2, dropout=0.00, use_qk_norm=True):
+        super().__init__()
+        self.d_model = d_model
+        self.d_output = d_output
+        self.half_dim = d_model // 2
+
+        # Dimensione delle query/key. Per efficienza la tengo pari a half_dim
+        self.d_k = self.half_dim
+
+        # Q = X W_q
+        self.W_q = nn.Linear(self.half_dim, self.d_k, bias=False)
+
+        # K = W_k (Matrice di parametri appresi, forma: [M, d_k])
+        self.W_k = nn.Parameter(torch.empty(M, self.d_k))
+        torch.nn.init.normal_(self.W_k, std=self.d_k ** -0.5)
+
+        # V = W_v corrisponde al tuo 'contract'
+        self.contract = nn.Linear(M, d_output, bias=False)
+
+        self.activation = activation()
+        self.dropout = nn.Dropout(dropout)
+
+        # QK Norm
+        self.use_qk_norm = use_qk_norm
+        if self.use_qk_norm:
+            self.q_norm = nn.RMSNorm(self.d_k)
+            self.k_norm = nn.RMSNorm(self.d_k)
+
+    def _init_weights(self):
+        torch.nn.init.zeros_(self.contract.weight)
+
+    def forward(self, x):
+        # 1. Prendiamo solo mezza residual stream
+        x_half = x[..., :self.half_dim]
+
+        # 2. Calcoliamo le Query e prendiamo le Key
+        q = self.W_q(x_half)
+        k = self.W_k
+
+        # 3. QK Norm
+        if self.use_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        # 4. Score dell'attention: Q K^T
+        # q è [..., d_k], k è [M, d_k]. F.linear(q, k) fa esattamente q @ k.T
+        scores = F.linear(q, k)
+
+        # 5. Attivazione generica (es. ReLU^2) e Dropout
+        h = self.dropout(self.activation(scores))
+
+        # 6. Proiezione V (contract) e padding automatico se d_output < d_model
+        out = self.contract(h)
+
+        if self.d_output < self.d_model:
+            out = F.pad(out, (0, self.d_model - self.d_output))
+
+        return out
+
+
+class VariableUDLP_CrossGate(Layer):
+    def __init__(self, d_model, d_output, M, activation=ReLU2, dropout=0.00):
+        super().__init__()
+        self.d_model = d_model
+        self.d_output = d_output
+        self.half_dim = d_model // 2
+
+        self.expand_k = nn.Linear(self.half_dim, M, bias=False)
+
+        self.expand_g = nn.Linear(d_model - self.half_dim, M, bias=False)
+
+        self.contract = nn.Linear(M, d_output, bias=False)
+
+        self.activation = activation()
+        self.dropout = nn.Dropout(dropout)
+
+    def _init_weights(self):
+        torch.nn.init.zeros_(self.contract.weight)
+
+    def forward(self, x):
+        x_left = x[..., :self.half_dim]
+        x_right = x[..., self.half_dim:]
+
+        hk = self.expand_k(x_left)
+        hg = self.expand_g(x_right)
+
+        h = self.dropout(self.activation(hk) * hg)
+
+        out = self.contract(h)
+
+        return out
+
+
+class VariableGatedUDLP(Layer):
+    """UDLP + sparse output gate (parameter-golf / Modded-NanoGPT 1667 stile).
+
+    Gate sigmoid su prima slice di `gate_input_dim` del residual (sparse).
+    Output: `contract(activation(expand(x))) * sigmoid(gate_proj(x[..., :gate_input_dim]))`.
+    """
+    __name__ = "VariableGatedUDLP"
+
+    def __init__(self, d_model, d_output, M, activation=ReLU2, dropout=0.00, gate_input_dim=12):
+        super().__init__()
+        self.gate_input_dim = gate_input_dim
+        self.expand = nn.Linear(d_model, M, bias=False)
+        self.contract = nn.Linear(M, d_output, bias=False)
+        self.gate_proj = nn.Linear(gate_input_dim, d_output, bias=False)
+        self.activation = activation()
+        self.dropout = nn.Dropout(dropout)
+
+    def _init_weights(self):
+        # Identity-at-init: expand + gate_proj orthogonal, contract → 0
+        # (con contract=0 il branch è zero indipendentemente dal gate).
+        torch.nn.init.orthogonal_(self.expand.weight)
+        torch.nn.init.zeros_(self.contract.weight)
+        torch.nn.init.orthogonal_(self.gate_proj.weight)
+
+    def forward(self, x):
+        out = self.contract(self.activation(self.expand(x)))
+        gate = torch.sigmoid(self.gate_proj(x[..., :self.gate_input_dim]))
+        return self.dropout(out * gate)
+
+
+class DoubleUDLP(Layer):
+    def __init__(self, d_model, d_output, M, activation=ReLU2, dropout=0.00):
+        super().__init__()
+        self.expand = DoubleLinear(d_model, M)
+        self.contract = nn.Linear(M, d_output)
+        self.activation = activation()
+        self.dropout = nn.Dropout(dropout)
+
+    def _init_weights(self):
+        torch.nn.init.zeros_(self.contract.weight)
+
+    def forward(self, x):
+        return self.dropout(self.contract(self.activation(self.expand(x))))
+
+
+class M_UDLP(Layer):
+    def __init__(self, d_model, d_output, M, activation=ReLU2, dropout=0.00):
+        super().__init__()
+        self.in_proj = nn.Linear(d_model, d_model)
+        self.expand = nn.Linear(d_model, M)
+        self.contract = nn.Linear(M, d_output)
+        self.activation = activation()
+        self.dropout = nn.Dropout(dropout)
+
+    def _init_weights(self):
+        torch.nn.init.zeros_(self.contract.weight)
+
+    def forward(self, x):
+        return self.dropout(self.contract(self.activation(self.expand(x))))
+
+
+class F_UDLPSwiGLU(Layer):
+    def __init__(self, d_model, expand, dropout=0.05, lora_rank=16):
+        super().__init__()
+        self.expand = nn.Linear(d_model, d_model * expand * 2, bias=False)
+        self.contract = nn.Linear(d_model * expand, d_model, bias=False)
+        self.LoRA_expand = LowRankLinear(d_model, expand * d_model * 2, lora_rank)
+        self.LoRA_contract = LowRankLinear(d_model * expand, d_model, lora_rank)
+        self.scale = 1 / lora_rank * 2
+        self.finetuning = False
+        self.activation = SwiGLU()
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        if not self.finetuning:
+            return self.dropout(self.contract(self.activation(self.expand(x))))
+        else:
+            act = self.activation(self.expand(x) + self.LoRA_expand(x) * self.scale)
+            return self.dropout(
+                self.contract(act) + self.LoRA_contract(act) * self.scale
+            )
+
+    def finetune(self):
+        torch.nn.init.zeros_(self.LORA_expand.l2.weight)
+        torch.nn.init.zeros_(self.LoRA_contract.l2.weight)
+        self.expand.requires_grad_(False)
+        self.contract.requires_grad_(False)
+        self.finetuning = True
+
+
+class ResMLPBlock(Layer):
+    def __init__(self, d_model, expand=2, norm=True, activation: Callable = nn.GELU):
+        super().__init__()
+        self.activation1 = activation()
+        self.activation2 = activation()
+        self.norm = norm
+        self.d_model = d_model
+        self.expand = expand
+        self.l1 = nn.Linear(d_model, d_model * expand, bias=True)
+        self.l2 = nn.Linear(d_model * expand, d_model, bias=True)
+
+        self.g1 = nn.Linear(d_model, d_model * expand, bias=True)
+        self.g2 = nn.Linear(d_model * expand, d_model, bias=True)
+
+        self.norm = nn.LayerNorm(d_model) if self.norm else nn.Identity()
+
+    def forward(self, x: torch.Tensor):
+        x = self.l2(self.activation1(self.l1(x)))
+        x = self.norm(x) + x
+        x = self.g2(self.activation2(self.g1(x)))
+        return x
+
+
+class ResMLP(Layer):
+    def __init__(self, d_model: int, depth: int, expand: int, activation: Callable, dropout=0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.depth = depth
+        self.expand = expand
+        self.activation = activation
+        self.dropout = nn.Dropout(dropout)
+
+        layers = []
+        for i in range(depth):
+            layers.append(ResMLPBlock(d_model, expand=expand, activation=activation))
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor):
+        return self.dropout(self.layers(x))
+
+
+class ConvResBlock(Layer):
+    __complexity__ = "O(L k^2 in out)"
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, activation=nn.ReLU):
+        super(ConvResBlock, self).__init__()
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size, padding='same')
+        self.conv2 = nn.Conv2d(in_channels, in_channels, kernel_size, padding='same')
+        self.convout = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
+        self.activation = activation()
+
+    def forward(self, x):
+        res = x
+        x = self.bn(x)
+        x = self.activation(self.conv1(x))
+        x = self.conv2(x)
+        x = x + res
+        return self.convout(x)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x):
+        input_dtype = x.dtype
+
+        x_fp32 = x.to(torch.float32)
+        variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
+
+        x_rsqrt = torch.rsqrt(variance + self.eps)
+
+        return self.weight * (x_fp32 * x_rsqrt).to(input_dtype)
+
+
+class EMA:
+    def __init__(self, model, decay=0.999):
+        self.decay = decay
+        self.shadow = {name: param.clone().detach() for name, param in model.named_parameters() if param.requires_grad}
+
+    @torch.no_grad()
+    def update(self, model):
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                self.shadow[name].lerp_(param, 1.0 - self.decay)
+
+    def apply_shadow(self, model):
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                param.data.copy_(self.shadow[name])
+
+
+class Nova(Layer):
+    def __init__(self):
+        super().__init__()
+        self.beta = nn.Parameter(torch.tensor([0.7]))
+
+    def forward(self, x):
+        h = -self.beta * x
+        return x * (torch.sigmoid(-h) - (1 / (1 + h ** 2)))
+
+
+class MinusNova(Layer):
+    def __init__(self):
+        super().__init__()
+        self.beta = nn.Parameter(torch.tensor([0.7]))
+
+    def forward(self, x):
+        h = -self.beta * -x
+        return x * (torch.sigmoid(-h) - (1 / (1 + h ** 2)))
+
+
+class TaylorAct(nn.Module):
+    def __init__(self, order=3):
+        super().__init__()
+        self.order = order
+        self.coeffs = nn.Parameter(torch.randn(order + 1) * 0.02)
+
+    def forward(self, x):
+        device = x.device
+        exponents = torch.arange(self.order + 1, device=device, dtype=x.dtype)
+        x_pow = x.unsqueeze(-1).pow(exponents)
+
+        return torch.matmul(x_pow, self.coeffs)
+
+    def plot(self, bounds=(-4, 4), n_points=200):
+        x = torch.linspace(bounds[0], bounds[1], n_points)
+        with torch.no_grad():
+            y = self.forward(x)
+
+        plt.figure(figsize=(7, 4))
+        plt.plot(x.numpy(), y.numpy(), label='TaylorAct', color='royalblue')
+        plt.axhline(0, color='gray', linewidth=0.8, linestyle='--')
+        plt.axvline(0, color='gray', linewidth=0.8, linestyle='--')
+        plt.title('TaylorAct Activation Function')
+        plt.xlabel('x')
+        plt.ylabel('f(x)')
+        plt.legend()
+        plt.grid(True, alpha=0.3)
+        plt.tight_layout()
+        plt.show()
+
+
+class FastPoly3(nn.Module):
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.a0 = nn.Parameter(torch.zeros(d_model))
+        self.a1 = nn.Parameter(torch.ones(d_model))
+        self.a2 = nn.Parameter(torch.zeros(d_model))
+        self.a3 = nn.Parameter(torch.zeros(d_model))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.a0 + x * (self.a1 + x * (self.a2 + x * self.a3))
+
+
+class CastedLinear(nn.Linear):
+    def forward(self, x: Tensor) -> Tensor:
+        bias = self.bias.to(x.dtype) if self.bias is not None else None
+        return F.linear(x, self.weight.to(x.dtype), bias)
+
+
+if __name__ == '__main__':
+    act = TaylorAct()
+    act.plot()  # default -4, 4
+    act.plot(bounds=(-1, 1))  # custom bounds
+    act.plot(bounds=(0, 1), n_points=100)

vathos/_experiments.py

@@ -0,0 +1,155 @@
+import json
+import os
+from pathlib import Path
+from datetime import datetime
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+
+
+class Experiment:
+    """
+    Un tracker di esperimenti stile WandB per il framework Vathos.
+    Gestisce configurazioni, metriche, salvataggi e visualizzazioni avanzate.
+    """
+
+    def __init__(self, project_name: str, exp_name: str = None, config: dict = None, root_dir: str = "runs"):
+        self.project_name = project_name
+
+        # Genera un nome univoco se non fornito
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.exp_name = exp_name if exp_name else f"run_{timestamp}"
+        self.config = config or {}
+
+        # Setup directory dell'esperimento
+        self.run_dir = Path(root_dir) / self.project_name / self.exp_name
+        self.run_dir.mkdir(parents=True, exist_ok=True)
+
+        # Dizionari per le metriche.
+        # Formato: {"loss": {"steps": [1,2,3], "values": [0.9, 0.8, 0.7]}}
+        self.metrics = {}
+
+        # Salva subito la configurazione
+        self._save_config()
+        print(f"🚀 Esperimento '{self.exp_name}' inizializzato in {self.run_dir}")
+
+    def log(self, metrics_dict: dict, step: int):
+        """
+        Registra un dizionario di metriche per un dato step.
+        Es: exp.log({"train_loss": 0.5, "val_accuracy": 0.8}, step=100)
+        """
+        for k, v in metrics_dict.items():
+            # Converte tensori PyTorch o Numpy in float standard
+            if hasattr(v, "item"):
+                v = v.item()
+
+            if k not in self.metrics:
+                self.metrics[k] = {"steps": [], "values": []}
+
+            self.metrics[k]["steps"].append(step)
+            self.metrics[k]["values"].append(v)
+
+    def _save_config(self):
+        with open(self.run_dir / "config.json", "w") as f:
+            json.dump(self.config, f, indent=4)
+
+    def save(self):
+        """Salva tutte le metriche correnti in un file JSON."""
+        with open(self.run_dir / "metrics.json", "w") as f:
+            json.dump(self.metrics, f, indent=4)
+        # print(f"💾 Metriche salvate in {self.run_dir / 'metrics.json'}")
+
+    @classmethod
+    def load(cls, run_dir: str):
+        """
+        Inizializza un oggetto Experiment a partire da una cartella salvata.
+        Permette di ricaricare e plottare run vecchie!
+        """
+        run_path = Path(run_dir)
+
+        # Estrai project e exp name dal path (es. runs/Progetto/Run)
+        exp_name = run_path.name
+        project_name = run_path.parent.name
+
+        with open(run_path / "config.json", "r") as f:
+            config = json.load(f)
+
+        # Ricostruisce l'oggetto bypassando la creazione della cartella root
+        exp = cls(project_name=project_name, exp_name=exp_name, config=config, root_dir=run_path.parent.parent)
+
+        with open(run_path / "metrics.json", "r") as f:
+            exp.metrics = json.load(f)
+
+        print(f"📂 Esperimento '{exp_name}' caricato con successo.")
+        return exp
+
+    def _smooth_curve(self, points, factor=0.85):
+        """Applica un Exponential Moving Average (EMA) per rendere i plot bellissimi."""
+        smoothed_points = []
+        for point in points:
+            if smoothed_points:
+                previous = smoothed_points[-1]
+                smoothed_points.append(previous * factor + point * (1 - factor))
+            else:
+                smoothed_points.append(point)
+        return smoothed_points
+
+    def plot(self, smoothing: float = 0.85, save_fig: bool = False, figsize=(10, 6)):
+        """
+        Genera plot in stile accademico/WandB.
+        Mostra le curve originali in trasparenza e le curve smoothate in evidenza.
+        """
+        if not self.metrics:
+            print("Nessuna metrica da plottare!")
+            return
+
+        # Stile estetico "stupendo" (simile a Seaborn/WandB)
+        plt.style.use('bmh')  # Usa un tema base pulito integrato in matplotlib
+
+        num_metrics = len(self.metrics)
+        cols = 2 if num_metrics > 1 else 1
+        rows = (num_metrics + 1) // 2
+
+        fig, axes = plt.subplots(rows, cols, figsize=(figsize[0] * cols, figsize[1] * rows))
+        if num_metrics == 1:
+            axes = [axes]
+        else:
+            axes = axes.flatten()
+
+        # Colori moderni
+        colors = list(mcolors.TABLEAU_COLORS.values())
+
+        for idx, (metric_name, data) in enumerate(self.metrics.items()):
+            ax = axes[idx]
+            steps = data["steps"]
+            values = data["values"]
+            color = colors[idx % len(colors)]
+
+            # Plotta i dati raw (sbiaditi)
+            ax.plot(steps, values, color=color, alpha=0.2, label='Raw')
+
+            # Plotta i dati smoothati (in evidenza)
+            if smoothing > 0 and len(values) > 5:
+                smoothed = self._smooth_curve(values, factor=smoothing)
+                ax.plot(steps, smoothed, color=color, linewidth=2.5, label=f'Smoothed (α={smoothing})')
+
+            ax.set_title(metric_name.replace("_", " ").title(), fontsize=14, fontweight='bold')
+            ax.set_xlabel("Steps", fontsize=11)
+            ax.set_ylabel("Value", fontsize=11)
+            ax.grid(True, linestyle='--', alpha=0.6)
+            ax.legend(loc="best", frameon=True, shadow=True)
+
+        for idx in range(num_metrics, len(axes)):
+            fig.delaxes(axes[idx])
+
+        plt.suptitle(f"Experiment: {self.project_name} - {self.exp_name}", fontsize=16, y=1.02)
+        plt.tight_layout()
+
+        if save_fig:
+            fig_path = self.run_dir / "plots.png"
+            plt.savefig(fig_path, dpi=300, bbox_inches='tight')
+            print(f"📊 Plot salvato in {fig_path}")
+
+        plt.show()
+
+

vathos/_spatials.py

@@ -0,0 +1,2124 @@
+from Vathos._basics import *
+
+
+class SinusoidalPositionalEncoding(Layer):
+    __name__ = "SinusoidalPositionalEncoding"
+    __complexity__ = "O(L^2 d^2)"
+
+    def __init__(self, d_model: int, max_len: int = 5000):
+        super().__init__()
+        self.d_model = d_model
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+
+        pe = torch.zeros(max_len, d_model)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        self.register_buffer('pe', pe)
+
+    def forward(self, x: torch.Tensor):
+        B, L, D = x.shape
+        return x + self.pe[:L]
+
+
+class RoPE(Layer):
+    __name__ = "RoPE"
+
+    def __init__(self, dim: int, max_len: int = 8192, base: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.base = base
+        self.max_len = max_len
+
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        self._cos_cached = None
+        self._sin_cached = None
+        self._seq_len_cached = 0
+
+    def _update_cache(self, seq_len: int, dtype: torch.dtype, device: torch.device):
+        if seq_len > self._seq_len_cached or self._cos_cached is None:
+            self._seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=device, dtype=dtype)
+            freqs = torch.outer(t, self.inv_freq.to(device))
+
+            self._cos_cached = freqs.cos()
+            self._sin_cached = freqs.sin()
+
+    def _apply_rotary_emb(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+                          start_pos: int = 0) -> torch.Tensor:
+        """Apply rotary embeddings starting from start_pos"""
+        seq_len = x.shape[-2]
+
+        cos = cos[start_pos:start_pos + seq_len]
+        sin = sin[start_pos:start_pos + seq_len]
+
+        shape = [1] * x.ndim
+        shape[-2] = seq_len
+        shape[-1] = self.dim // 2
+
+        cos = cos.view(*shape)
+        sin = sin.view(*shape)
+
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor = None, start_pos: int = 0):
+        """
+        Args:
+            q: Query tensor
+            k: Key tensor (optional)
+            start_pos: Starting position for RoPE (used during generation)
+        """
+        assert q.shape[-1] == self.dim, f"Last dim of q must be {self.dim}, got {q.shape[-1]}"
+        if k is not None:
+            assert k.shape[-2:] == q.shape[-2:], "k must have same seq_len and head_dim as q"
+
+        seq_len = q.shape[-2]
+        self._update_cache(start_pos + seq_len, q.dtype, q.device)
+
+        cos = self._cos_cached
+        sin = self._sin_cached
+
+        q_rope = self._apply_rotary_emb(q, cos, sin, start_pos)
+        k_rope = self._apply_rotary_emb(k, cos, sin, start_pos) if k is not None else None
+
+        return (q_rope, k_rope) if k_rope is not None else q_rope
+
+
+class ALiBi(Layer):
+    __name__ = "ALiBi"
+
+    def __init__(self, n_heads: int):
+        super().__init__()
+        self.n_heads = n_heads
+
+        # Calculate standard ALiBi geometric slopes
+        closest_power_of_2 = 2 ** math.floor(math.log2(n_heads))
+        base = 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3)))
+        slopes = [math.pow(base, i) for i in range(1, closest_power_of_2 + 1)]
+
+        # Handle non-power-of-2 head counts (e.g., 12 heads)
+        if closest_power_of_2 != n_heads:
+            extra_base = 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3)))
+            slopes.extend([math.pow(extra_base, i) for i in range(1, 2 * (n_heads - closest_power_of_2) + 1, 2)])
+
+        self.register_buffer(
+            "slopes",
+            torch.tensor(slopes, dtype=torch.float32).view(n_heads, 1, 1),
+            persistent=False
+        )
+
+    def forward(self, q_len: int, k_len: int, start_pos: int, device: torch.device, dtype: torch.dtype,
+                causal: bool = True) -> torch.Tensor:
+        # Create position indices
+        q_idx = torch.arange(start_pos, start_pos + q_len, device=device)[:, None]
+        k_idx = torch.arange(0, k_len, device=device)[None, :]
+
+        # Calculate relative distances (j - i).
+        # For causal attention, we only attend to the past, so k_idx <= q_idx, meaning distances <= 0.
+        distances = k_idx - q_idx
+
+        # Scale distances by head-specific slopes
+        alibi_bias = distances * self.slopes
+        alibi_bias = alibi_bias.to(dtype)
+
+        # Embed the causal mask directly into the ALiBi bias
+        if causal:
+            causal_mask = distances > 0
+            alibi_bias.masked_fill_(causal_mask, float('-inf'))
+
+        # Add batch dimension for SDPA broadcasting -> (1, n_heads, q_len, k_len)
+        return alibi_bias.unsqueeze(0)
+
+
+class YaRN(Layer):
+    __name__ = "YaRN"
+
+    def __init__(self, dim: int, original_max_len: int = 4096, scale: float = 1.0,
+                 base: float = 10000.0, beta_fast: int = 32, beta_slow: int = 1):
+        super().__init__()
+        self.dim = dim
+        self.base = base
+        self.original_max_len = original_max_len
+        self.scale = scale
+
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+
+        if scale > 1.0:
+            wavelengths = 2 * math.pi / inv_freq
+
+            low = original_max_len / beta_slow
+            high = original_max_len / beta_fast
+
+            w = torch.clamp((wavelengths - high) / (low - high), 0.0, 1.0)
+
+            inv_freq_interpolated = inv_freq / scale
+            inv_freq_extrapolated = inv_freq
+
+            inv_freq = (1 - w) * inv_freq_extrapolated + w * inv_freq_interpolated
+
+            self.mscale = math.sqrt(0.1 * math.log(scale) + 1.0)
+        else:
+            self.mscale = 1.0
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        self._cos_cached = None
+        self._sin_cached = None
+        self._seq_len_cached = 0
+
+    def _update_cache(self, seq_len: int, dtype: torch.dtype, device: torch.device):
+        if seq_len > self._seq_len_cached or self._cos_cached is None:
+            self._seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=device, dtype=dtype)
+            freqs = torch.outer(t, self.inv_freq.to(device))
+
+            # Apply the YaRN mscale directly to the cache
+            self._cos_cached = (freqs.cos() * self.mscale)
+            self._sin_cached = (freqs.sin() * self.mscale)
+
+    def _apply_rotary_emb(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+                          start_pos: int = 0) -> torch.Tensor:
+        """Apply rotary embeddings starting from start_pos"""
+        seq_len = x.shape[-2]
+
+        cos = cos[start_pos:start_pos + seq_len]
+        sin = sin[start_pos:start_pos + seq_len]
+
+        shape = [1] * x.ndim
+        shape[-2] = seq_len
+        shape[-1] = self.dim // 2
+
+        cos = cos.view(*shape).to(x.dtype)
+        sin = sin.view(*shape).to(x.dtype)
+
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor = None, start_pos: int = 0):
+        assert q.shape[-1] == self.dim, f"Last dim of q must be {self.dim}, got {q.shape[-1]}"
+        if k is not None:
+            assert k.shape[-2:] == q.shape[-2:], "k must have same seq_len and head_dim as q"
+
+        seq_len = q.shape[-2]
+        self._update_cache(start_pos + seq_len, q.dtype, q.device)
+
+        cos = self._cos_cached
+        sin = self._sin_cached
+
+        q_rope = self._apply_rotary_emb(q, cos, sin, start_pos)
+        k_rope = self._apply_rotary_emb(k, cos, sin, start_pos) if k is not None else None
+
+        return (q_rope, k_rope) if k_rope is not None else q_rope
+
+
+class NanoYaRN(Layer):
+    """Modded Nano-Gpt implmenetation of YaRN"""
+
+    def __init__(self, head_dim, max_seq_len, paired=False):
+        super().__init__()
+        self.head_dim = head_dim
+        self.max_seq_len = max_seq_len
+        self.paired = paired
+        self.reset()
+
+    def rotary(self, x_BTHD):
+        assert self.factor1.size(0) >= x_BTHD.size(-3)
+        factor1, factor2 = (
+            self.factor1[None, : x_BTHD.size(-3), None, :],
+            self.factor2[None, : x_BTHD.size(-3), None, :],
+        )
+        x_flip = x_BTHD.view(*x_BTHD.shape[:-1], x_BTHD.shape[-1] // 2, 2).flip(-1).view(x_BTHD.shape)
+        return factor1 * x_BTHD + factor2 * x_flip
+
+    def reset(self):
+        angular_freq = (1 / 1024) ** torch.linspace(0, 1, steps=self.head_dim // 4, dtype=torch.float32)
+        angular_freq = angular_freq.repeat_interleave(2)
+        # half-truncate RoPE by @YouJiacheng (w/ base freq tuning)
+        angular_freq = torch.cat([angular_freq, angular_freq.new_zeros(self.head_dim // 2)])
+        t = torch.arange(2 * self.max_seq_len, dtype=torch.float32)
+        if not self.paired:
+            theta = torch.outer(t, angular_freq)
+            self.factor1 = nn.Buffer(
+                theta.cos().to(torch.bfloat16), persistent=False
+            )
+            self.factor2 = nn.Buffer(
+                theta.sin().to(torch.bfloat16), persistent=False
+            )
+        else:
+            t_even = 2 * t
+            t_odd = 2 * t + 1
+            theta1 = torch.outer(t_even, angular_freq)
+            theta2 = torch.outer(t_odd, angular_freq)
+            self.factor1 = nn.Buffer(
+                torch.cat((theta1.cos(), theta2.cos()), dim=-1).to(torch.bfloat16),
+                persistent=False
+            )
+            self.factor2 = nn.Buffer(
+                torch.cat((theta1.sin(), theta2.sin()), dim=-1).to(torch.bfloat16),
+                persistent=False
+            )
+        self.factor2[..., 1::2] *= -1
+        self.angular_freq = angular_freq
+        # start with 0.1, inspired by 0.12 from @leloykun and learnable scalars used by @brendanh0gan https://x.com/hi_tysam/status/1879693583898591283
+        self.attn_scale = 0.1
+
+    def apply(self, old_window: int, new_window: int, alpha: int = 1, beta: int = 32):
+        rotations = old_window * self.angular_freq / (2 * torch.pi)
+        scaling_factor = old_window / new_window
+        interpolation_weight = torch.clamp((rotations - alpha) / (beta - alpha), 0, 1)
+        self.angular_freq *= scaling_factor + interpolation_weight * (1 - scaling_factor)
+        t = torch.arange(2 * self.max_seq_len, dtype=torch.float32, device=self.angular_freq.device)
+        if not self.paired:
+            theta = torch.outer(t, self.angular_freq)
+            self.factor1.copy_(theta.cos())
+            self.factor2.copy_(theta.sin())
+        else:
+            t_even = 2 * t
+            t_odd = 2 * t + 1
+            theta1 = torch.outer(t_even, self.angular_freq)
+            theta2 = torch.outer(t_odd, self.angular_freq)
+            self.factor1.copy_(torch.cat((theta1.cos(), theta2.cos()), dim=-1))
+            self.factor2.copy_(torch.cat((theta1.sin(), theta2.sin()), dim=-1))
+        self.factor2[..., 1::2] *= -1
+        self.attn_scale *= 0.2 * math.log(new_window / old_window) + 1
+
+
+class MultiheadAttentionMixer(Layer):
+    __name__ = "MultiheadAttentionMixer"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool = True,
+                 pos_emb: nn.Module = None, dropout: float = 0.0, qk_norm=False):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        # Identity-at-init: qkv orthogonal (preserva geometria), out → 0.
+        nn.init.orthogonal_(self.qkv.weight)
+        nn.init.zeros_(self.out.weight)
+
+    # Alias backward-compat per chi chiama explicit _reset_parameters.
+    _reset_parameters = _init_weights
+
+    def forward(self, x: torch.Tensor, ve=None) -> torch.Tensor:
+        """
+        ve: Stands for value embeddings, which should be already weighted byt the caller, of course the dimension must match x
+        """
+        B, L, D = x.shape
+
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+        if ve is not None:
+            v = v + ve.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def show_weights_forward(self, x):
+        B, L, D = x.shape
+
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
+
+        if self.causal:
+            mask = torch.triu(torch.ones(L, L, dtype=torch.bool, device=x.device), diagonal=1)
+            scores.masked_fill_(mask, float('-inf'))
+
+        attn_weights = F.softmax(scores, dim=-1)
+        attn = torch.matmul(attn_weights, v)
+
+        out = self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+        return out, attn_weights
+
+    def generate(self, x: torch.Tensor, ve=None) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=pos_offset)
+
+        if ve is not None:
+            v = v + ve.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal and (L > 1),
+            dropout_p=0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+    def finetune(self):
+        self.qkv.weight.requires_grad = False
+        self.out.weight.requires_grad = False
+
+
+class MultiheadGatedAttentionMixer(Layer):
+    """MHA + sparse per-head output gate (Modded-NanoGPT PR #117 / parameter-golf 1667).
+
+    Gate input: prime `gate_input_dim` (default 12) dimensioni di x — sparse, ottimizzato
+    per Muon. Output: un sigmoid scalare per head, broadcast su head_dim. Init zero ⇒
+    gate parte a 0.5 (no-op).
+    """
+    __name__ = "MultiheadGatedAttentionMixer"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool = True,
+                 pos_emb: nn.Module = None, dropout: float = 0.0, qk_norm: bool = False,
+                 gate_input_dim: int = 12):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+        self.gate_input_dim = gate_input_dim
+
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+        self.gate_proj = nn.Linear(gate_input_dim, n_heads, bias=False)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+        self.qk_norm = qk_norm
+        if qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        # Identity-at-init: qkv orthogonal, out → 0, gate_proj orthogonal
+        # (out=0 ⇒ contributo attention nullo all'init, indipendentemente dal gate).
+        nn.init.orthogonal_(self.qkv.weight)
+        nn.init.zeros_(self.out.weight)
+        nn.init.orthogonal_(self.gate_proj.weight)
+
+    _reset_parameters = _init_weights
+
+    def _apply_gate(self, attn: torch.Tensor, x_gate_in: torch.Tensor) -> torch.Tensor:
+        # attn: [B, H, L, D_h] | x_gate_in: [B, L, gate_input_dim]
+        gate = torch.sigmoid(self.gate_proj(x_gate_in))                 # [B, L, H]
+        return attn * gate.transpose(1, 2).unsqueeze(-1)                # broadcast su D_h
+
+    def forward(self, x: torch.Tensor, ve=None) -> torch.Tensor:
+        B, L, D = x.shape
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+        if ve is not None:
+            v = v + ve.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v, is_causal=self.causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+        attn = self._apply_gate(attn, x[..., :self.gate_input_dim])
+        return self.out(attn.transpose(1, 2).reshape(B, L, D))
+
+    def generate(self, x: torch.Tensor, ve=None) -> torch.Tensor:
+        B, L, D = x.shape
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=pos_offset)
+        if ve is not None:
+            v = v + ve.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+        self.kv_cache = (k, v)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v, is_causal=self.causal and (L > 1), dropout_p=0.0,
+        )
+        attn = self._apply_gate(attn, x[..., :self.gate_input_dim])
+        return self.out(attn.transpose(1, 2).reshape(B, L, D))
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+
+class MultiheadAttentionMixerALIBI(Layer):
+    __name__ = "MultiheadAttentionMixer"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool = True,
+                 pos_emb: nn.Module = None, alibi: nn.Module = None, dropout: float = 0.0, qk_norm=False):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+
+        self.pos_emb = pos_emb
+        self.alibi = alibi  # ADDED: ALiBi module injection
+        self.kv_cache = None
+        self.qk_norm = qk_norm
+
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.qkv.weight)
+        nn.init.xavier_uniform_(self.out.weight)
+
+    def forward(self, x: torch.Tensor, ve=None) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        if ve is not None:
+            v = v + ve.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        # ADDED: ALiBi Mask Logic
+        attn_mask = None
+        is_causal = self.causal
+        if self.alibi is not None:
+            attn_mask = self.alibi(q_len=L, k_len=L, start_pos=0, device=x.device, dtype=q.dtype, causal=self.causal)
+            is_causal = False  # PyTorch requires is_causal=False if a custom attn_mask is provided
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def show_weights_forward(self, x):
+        B, L, D = x.shape
+
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
+
+        # ADDED: ALiBi logic replaces strict causal mask if active
+        if self.alibi is not None:
+            alibi_bias = self.alibi(q_len=L, k_len=L, start_pos=0, device=x.device, dtype=q.dtype, causal=self.causal)
+            scores = scores + alibi_bias
+        elif self.causal:
+            mask = torch.triu(torch.ones(L, L, dtype=torch.bool, device=x.device), diagonal=1)
+            scores.masked_fill_(mask, float('-inf'))
+
+        attn_weights = F.softmax(scores, dim=-1)
+        attn = torch.matmul(attn_weights, v)
+
+        out = self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+        return out, attn_weights
+
+    def generate(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=pos_offset)
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+        k_len = k.shape[2]
+
+        # ADDED: ALiBi Mask Logic for Generation
+        attn_mask = None
+        is_causal = self.causal and (L > 1)
+        if self.alibi is not None:
+            attn_mask = self.alibi(q_len=L, k_len=k_len, start_pos=pos_offset, device=x.device, dtype=q.dtype,
+                                   causal=self.causal)
+            is_causal = False
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+            dropout_p=0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    # [clear_cache and finetune remain unchanged]
+
+
+class MultiheadAttentionMixerXSA(Layer):
+    __name__ = "MultiheadAttentionMixerXSA"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool = True,
+                 pos_emb: nn.Module = None, dropout: float = 0.0, qk_norm=False):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.qkv.weight)
+        nn.init.xavier_uniform_(self.out.weight)
+
+    def forward(self, x: torch.Tensor, ve=None) -> torch.Tensor:
+        """
+        ve: Stands for value embeddings, which should be already weighted by the caller, of course the dimension must match x
+        """
+        B, L, D = x.shape
+
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+        if ve is not None:
+            v = v + ve
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+
+        v_norm = F.normalize(v, dim=-1)
+        attn = attn - (attn * v_norm).sum(dim=-1, keepdim=True) * v_norm
+
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def show_weights_forward(self, x):
+        B, L, D = x.shape
+
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
+
+        if self.causal:
+            mask = torch.triu(torch.ones(L, L, dtype=torch.bool, device=x.device), diagonal=1)
+            scores.masked_fill_(mask, float('-inf'))
+
+        attn_weights = F.softmax(scores, dim=-1)
+        attn = torch.matmul(attn_weights, v)
+
+        v_norm = F.normalize(v, dim=-1)
+        attn = attn - (attn * v_norm).sum(dim=-1, keepdim=True) * v_norm
+
+        out = self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+        return out, attn_weights
+
+    def generate(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qkv = self.qkv(x).view(B, L, 3, self.n_heads, self.head_dim)
+        q, k, v = qkv.unbind(dim=2)
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=pos_offset)
+
+        v_self = v
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal and (L > 1),
+            dropout_p=0.0,
+        )
+
+        v_norm = F.normalize(v_self, dim=-1)
+        attn = attn - (attn * v_norm).sum(dim=-1, keepdim=True) * v_norm
+
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+    def finetune(self):
+        self.qkv.weight.requires_grad = False
+        self.out.weight.requires_grad = False
+
+
+class MultiheadAttentionMixerNOV(Layer):
+    __name__ = "MultiheadAttentionMixerNOV"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool = True,
+                 pos_emb: nn.Module = None, dropout: float = 0.0, qk_norm=False):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+
+        self.qk = nn.Linear(d_model, 2 * d_model, bias=False)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.qk.weight)
+        nn.init.xavier_uniform_(self.out.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qk = self.qk(x).view(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.unbind(dim=2)
+        q, k = q.transpose(1, 2), k.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        v = x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def generate(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qk = self.qk(x).view(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.unbind(dim=2)
+        q, k = q.transpose(1, 2), k.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        v = x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=pos_offset)
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal and (L > 1),
+            dropout_p=0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+    def finetune(self):
+        self.qk.weight.requires_grad = False
+        self.out.weight.requires_grad = False
+
+
+class MultiheadAttentionMixerNOVLa2(Layer):
+    __name__ = "MultiheadAttentionMixerNOV"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool = True,
+                 pos_emb: nn.Module = None, dropout: float = 0.0, qk_norm=False):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+
+        self.qk = nn.Linear(d_model, 2 * d_model, bias=False)
+        self.sec = nn.Linear(d_model, d_model, bias=False)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.xavier_uniform_(self.qk.weight)
+        nn.init.xavier_uniform_(self.out.weight)
+        nn.init.xavier_uniform_(self.sec.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qk = self.qk(self.sec(x)).view(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.unbind(dim=2)
+        q, k = q.transpose(1, 2), k.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        v = x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def generate(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qk = self.qk(x).view(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.unbind(dim=2)
+        q, k = q.transpose(1, 2), k.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        v = x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=pos_offset)
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal and (L > 1),
+            dropout_p=0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+    def finetune(self):
+        self.qk.weight.requires_grad = False
+        self.out.weight.requires_grad = False
+
+
+class MultiheadAttentionMixerNOVLa3(Layer):
+    __name__ = "MultiheadAttentionMixerNOV"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool = True,
+                 pos_emb: nn.Module = None, dropout: float = 0.0, qk_norm=False):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+
+        self.qk = DoubleLinear(d_model, 2 * d_model)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.xavier_uniform_(self.out.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qk = self.qk(x).view(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.unbind(dim=2)
+        q, k = q.transpose(1, 2), k.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        v = x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def generate(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qk = self.qk(x).view(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.unbind(dim=2)
+        q, k = q.transpose(1, 2), k.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        v = x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=pos_offset)
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal and (L > 1),
+            dropout_p=0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+    def finetune(self):
+        self.qk.weight.requires_grad = False
+        self.out.weight.requires_grad = False
+
+
+class MultiheadAttentionMixerNOVLa4(Layer):
+    __name__ = "MultiheadAttentionMixerNOV"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool = True,
+                 pos_emb: nn.Module = None, dropout: float = 0.0, qk_norm=False):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+
+        self.qk = DoubleLinear(d_model, 2 * d_model)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.out.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qk = self.qk(x).view(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.unbind(dim=2)
+        q, k = q.transpose(1, 2), k.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        v = x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def generate(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qk = self.qk(x).view(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.unbind(dim=2)
+        q, k = q.transpose(1, 2), k.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        v = x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=pos_offset)
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal and (L > 1),
+            dropout_p=0.0,
+        )
+        return self.out(attn.transpose(1, 2).contiguous().view(B, L, D))
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+    def finetune(self):
+        self.qk.weight.requires_grad = False
+        self.out.weight.requires_grad = False
+
+
+class GQAGatedMixer(Layer):
+    """GQA + sparse per-head output gate + VE + KV cache.
+
+    Combina:
+    - GQA: n_kv_heads << n_heads. KV cache compatto (n_kv_heads heads).
+    - Sparse output gate: sigmoid(gate_proj(x[..., :gate_input_dim])) per-head,
+      broadcast su head_dim (Modded-NanoGPT PR #117 stile).
+    - Value embeddings (ve): aggiunti a v_expanded (full n_heads). Quando ve
+      è attivo la cache memorizza v in formato **expanded** (n_heads, costa
+      n_rep× memoria) per preservare per-step ve correttamente. Quando ve è
+      None, cache in formato compresso standard (n_kv_heads).
+
+    Identity-at-init: q_proj/kv_proj/gate_proj orthogonal, o_proj → 0.
+    """
+    __name__ = "GQAGatedMixer"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int,
+                 causal: bool = True, pos_emb: nn.Module = None,
+                 dropout: float = 0.0, qk_norm: bool = False,
+                 gate_input_dim: int = 12):
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError(f"d_model ({d_model}) must be divisible by n_heads ({n_heads})")
+        if n_heads % n_kv_heads != 0:
+            raise ValueError(f"n_heads ({n_heads}) must be divisible by n_kv_heads ({n_kv_heads})")
+
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = d_model // n_heads
+        self.n_rep = n_heads // n_kv_heads
+        self.causal = causal
+        self.dropout_p = dropout
+        self.gate_input_dim = gate_input_dim
+
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.kv_proj = nn.Linear(d_model, n_kv_heads * self.head_dim * 2, bias=False)
+        self.o_proj = nn.Linear(d_model, d_model, bias=False)
+        self.gate_proj = nn.Linear(gate_input_dim, n_heads, bias=False)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+
+        self.qk_norm = qk_norm
+        if qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        # Identity-at-init
+        nn.init.orthogonal_(self.q_proj.weight)
+        nn.init.orthogonal_(self.kv_proj.weight)
+        nn.init.zeros_(self.o_proj.weight)
+        nn.init.orthogonal_(self.gate_proj.weight)
+
+    _reset_parameters = _init_weights
+
+    def _expand_kv(self, x: torch.Tensor) -> torch.Tensor:
+        if self.n_rep == 1:
+            return x
+        B, H, L, D = x.shape
+        return (x.unsqueeze(2)
+                .expand(B, H, self.n_rep, L, D)
+                .reshape(B, H * self.n_rep, L, D))
+
+    def _apply_gate(self, attn: torch.Tensor, x_gate_in: torch.Tensor) -> torch.Tensor:
+        # attn: [B, n_heads, L, head_dim] | x_gate_in: [B, L, gate_input_dim]
+        gate = torch.sigmoid(self.gate_proj(x_gate_in))                 # [B, L, H]
+        return attn * gate.transpose(1, 2).unsqueeze(-1)
+
+    def forward(self, x: torch.Tensor, ve: torch.Tensor = None) -> torch.Tensor:
+        B, T, D = x.shape
+
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        kv = self.kv_proj(x).view(B, T, self.n_kv_heads, 2, self.head_dim)
+        k, v = kv.unbind(dim=3)
+        k, v = k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        k_exp = self._expand_kv(k)
+        v_exp = self._expand_kv(v)
+        if ve is not None:
+            v_exp = v_exp + ve.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+
+        attn = F.scaled_dot_product_attention(
+            q, k_exp, v_exp, is_causal=self.causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+        attn = self._apply_gate(attn, x[..., :self.gate_input_dim])
+        return self.o_proj(attn.transpose(1, 2).reshape(B, T, D))
+
+    def generate(self, x: torch.Tensor, ve: torch.Tensor = None) -> torch.Tensor:
+        B, T, D = x.shape
+
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        kv = self.kv_proj(x).view(B, T, self.n_kv_heads, 2, self.head_dim)
+        k_new, v_new = kv.unbind(dim=3)
+        k_new, v_new = k_new.transpose(1, 2), v_new.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k_new = self.k_norm(k_new)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+        if self.pos_emb is not None:
+            q, k_new = self.pos_emb(q, k_new, start_pos=pos_offset)
+
+        # Cache k sempre in formato raw (n_kv_heads). Cache v:
+        # - se ve attivo: formato EXPANDED (n_heads), perché ve si applica per-step
+        # - altrimenti: raw (n_kv_heads), standard GQA
+        if ve is not None:
+            v_new = self._expand_kv(v_new) + ve.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k_new = torch.cat([k_cache, k_new], dim=2)
+            v_new = torch.cat([v_cache, v_new], dim=2)
+        self.kv_cache = (k_new, v_new)
+
+        # Attention: k_new sempre da espandere; v_new già expanded se ve, raw altrimenti
+        k_for_attn = self._expand_kv(k_new)
+        v_for_attn = v_new if ve is not None else self._expand_kv(v_new)
+
+        attn = F.scaled_dot_product_attention(
+            q, k_for_attn, v_for_attn,
+            is_causal=self.causal and (T > 1), dropout_p=0.0,
+        )
+        attn = self._apply_gate(attn, x[..., :self.gate_input_dim])
+        return self.o_proj(attn.transpose(1, 2).reshape(B, T, D))
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+
+class GroupedQueryAttention(Layer):
+    __name__ = "GroupedQueryAttention"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int,
+                 causal: bool = True, pos_emb: nn.Module = None,
+                 dropout: float = 0.0, bias: bool = False, qk_norm=False):
+        super().__init__()
+        if d_model % n_heads != 0:
+            raise ValueError(f"d_model ({d_model}) must be divisible by n_heads ({n_heads})")
+        if n_heads % n_kv_heads != 0:
+            raise ValueError(f"n_heads ({n_heads}) must be divisible by n_kv_heads ({n_kv_heads})")
+
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = d_model // n_heads
+        self.n_rep = n_heads // n_kv_heads
+        self.causal = causal
+        self.dropout_p = dropout
+
+        self.q_proj = nn.Linear(d_model, d_model, bias=bias)
+        self.kv_proj = nn.Linear(d_model, n_kv_heads * self.head_dim * 2, bias=bias)
+        self.o_proj = nn.Linear(d_model, d_model, bias=bias)
+
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.q_proj.weight)
+        nn.init.xavier_uniform_(self.kv_proj.weight)
+        nn.init.xavier_uniform_(self.o_proj.weight)
+        if self.q_proj.bias is not None:
+            nn.init.zeros_(self.q_proj.bias)
+            nn.init.zeros_(self.kv_proj.bias)
+            nn.init.zeros_(self.o_proj.bias)
+
+    def _expand_kv(self, x: torch.Tensor) -> torch.Tensor:
+        if self.n_rep == 1:
+            return x
+        B, H, L, D = x.shape
+        return (x.unsqueeze(2)
+                .expand(B, H, self.n_rep, L, D)
+                .reshape(B, H * self.n_rep, L, D))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+
+        kv = self.kv_proj(x).view(B, T, self.n_kv_heads, 2, self.head_dim)
+        k, v = kv.unbind(dim=3)
+        k, v = k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=0)
+
+        k, v = self._expand_kv(k), self._expand_kv(v)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal,
+            dropout_p=self.dropout_p if self.training else 0.0,
+        )
+        return self.o_proj(attn.transpose(1, 2).contiguous().view(B, T, C))
+
+    def generate(self, x: torch.Tensor, ve=None) -> torch.Tensor:
+        B, T, C = x.shape
+
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+
+        kv = self.kv_proj(x).view(B, T, self.n_kv_heads, 2, self.head_dim)
+        k, v = kv.unbind(dim=3)
+        k, v = k.transpose(1, 2), v.transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        pos_offset = self.kv_cache[0].shape[2] if self.kv_cache is not None else 0
+
+        if self.pos_emb is not None:
+            q, k = self.pos_emb(q, k, start_pos=pos_offset)
+
+        if ve is not None:
+            raise NotImplementedError(
+                "GroupedQueryAttention.forward does not support value embeddings; "
+                "use MultiheadAttentionMixer if you need ve."
+            )
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+
+        k, v = self._expand_kv(k), self._expand_kv(v)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=self.causal and (T > 1),
+            dropout_p=0.0,
+        )
+        return self.o_proj(attn.transpose(1, 2).contiguous().view(B, T, C))
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+    def finetune(self):
+        self.q_proj.weight.requires_grad = False
+        self.kv_proj.weight.requires_grad = False
+        self.o_proj.weight.requires_grad = False
+
+
+class GroupedQueryAttentionNOV(Layer):
+
+    def __init__(
+            self,
+            d_model: int,
+            n_heads: int,
+            n_kv_heads: int,
+            dropout: float = 0.0,
+            bias: bool = False,
+            causal: bool = True,
+            rope: bool = False,
+            qk_norm: bool = False
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.num_heads = n_heads
+        self.num_kv_heads = n_kv_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_prob = dropout
+        self.causal = causal
+
+        self.n_rep = self.num_heads // self.num_kv_heads
+
+        if self.d_model % self.num_heads != 0:
+            raise ValueError(f"embed_dim ({d_model}) must be divisible by num_heads ({n_heads})")
+        if self.num_heads % self.num_kv_heads != 0:
+            raise ValueError(f"num_heads ({n_heads}) must be divisible by num_kv_heads ({n_kv_heads})")
+
+        self.q_proj = nn.Linear(d_model, d_model, bias=bias)
+        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=bias)
+        self.o_proj = nn.Linear(d_model, d_model, bias=bias)
+
+        self.rope = RoPE(self.head_dim) if rope else None
+
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.q_proj.weight)
+        nn.init.xavier_uniform_(self.k_proj.weight)
+        nn.init.xavier_uniform_(self.o_proj.weight)
+        if self.q_proj.bias is not None:
+            nn.init.constant_(self.q_proj.bias, 0)
+            nn.init.constant_(self.k_proj.bias, 0)
+            nn.init.constant_(self.o_proj.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+
+        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = x.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        if self.rope is not None:
+            q, k = self.rope(q, k)
+
+        if self.n_rep > 1:
+            k = k[:, :, None, :, :].expand(B, self.num_kv_heads, self.n_rep, T, self.head_dim)
+            k = k.reshape(B, self.num_heads, T, self.head_dim)
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.dropout_prob if self.training else 0.0,
+            is_causal=self.causal
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, C)
+        return self.o_proj(attn_output)
+
+
+class GroupedQueryAttentionNOV2(Layer):
+    def __init__(
+            self,
+            d_model: int,
+            n_heads: int,
+            n_kv_heads: int,
+            dropout: float = 0.0,
+            bias: bool = False,
+            causal: bool = True,
+            rope: bool = False
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.num_heads = n_heads
+        self.num_kv_heads = n_kv_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_prob = dropout
+        self.causal = causal
+
+        self.n_rep = self.num_heads // self.num_kv_heads
+
+        if self.d_model % self.num_heads != 0:
+            raise ValueError(f"embed_dim ({d_model}) must be divisible by num_heads ({n_heads})")
+        if self.num_heads % self.num_kv_heads != 0:
+            raise ValueError(f"num_heads ({n_heads}) must be divisible by num_kv_heads ({n_kv_heads})")
+
+        self.q_proj = nn.Linear(d_model, d_model, bias=bias)
+        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=bias)
+        self.o_proj = nn.Linear(d_model, d_model, bias=bias)
+        self.trainv = False
+        self.v_proj = nn.Linear(d_model, d_model, bias=bias)
+
+        self.rope = RoPE(self.head_dim) if rope else None
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.q_proj.weight)
+        nn.init.xavier_uniform_(self.k_proj.weight)
+        nn.init.xavier_uniform_(self.o_proj.weight)
+        nn.init.eye_(self.v_proj.weight)
+        if self.q_proj.bias is not None:
+            nn.init.constant_(self.q_proj.bias, 0)
+            nn.init.constant_(self.k_proj.bias, 0)
+            nn.init.constant_(self.v_proj.bias, 0)
+            nn.init.constant_(self.o_proj.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+
+        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        if not self.trainv:
+            v = x.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        else:
+            v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Apply RoPE BEFORE expansion
+        if self.rope is not None:
+            q, k = self.rope(q, k)
+
+        if self.n_rep > 1:
+            k = k[:, :, None, :, :].expand(B, self.num_kv_heads, self.n_rep, T, self.head_dim)
+            k = k.reshape(B, self.num_heads, T, self.head_dim)
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.dropout_prob if self.training else 0.0,
+            is_causal=self.causal
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, C)
+        return self.o_proj(attn_output)
+
+
+class GroupedQueryAttentionNOO(Layer):
+    def __init__(
+            self,
+            d_model: int,
+            n_heads: int,
+            n_kv_heads: int,
+            dropout: float = 0.0,
+            bias: bool = False,
+            causal: bool = True,
+            rope: bool = False
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.num_heads = n_heads
+        self.num_kv_heads = n_kv_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_prob = dropout
+        self.causal = causal
+
+        self.n_rep = self.num_heads // self.num_kv_heads
+
+        if self.d_model % self.num_heads != 0:
+            raise ValueError(f"embed_dim ({d_model}) must be divisible by num_heads ({n_heads})")
+        if self.num_heads % self.num_kv_heads != 0:
+            raise ValueError(f"num_heads ({n_heads}) must be divisible by num_kv_heads ({n_kv_heads})")
+
+        self.q_proj = nn.Linear(d_model, d_model, bias=bias)
+        self.kv_proj = nn.Linear(d_model, n_kv_heads * self.head_dim * 2, bias=bias)
+
+        self.rope = RoPE(self.head_dim) if rope else None
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.q_proj.weight)
+        nn.init.xavier_uniform_(self.kv_proj.weight)
+        if self.q_proj.bias is not None:
+            nn.init.constant_(self.q_proj.bias, 0)
+            nn.init.constant_(self.kv_proj.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.shape
+
+        q = self.q_proj(x)
+        kv = self.kv_proj(x)
+
+        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        kv = kv.view(B, T, self.num_kv_heads, 2, self.head_dim)
+        k, v = kv.unbind(dim=3)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        # Apply RoPE BEFORE expansion
+        if self.rope is not None:
+            q, k = self.rope(q, k)
+
+        if self.n_rep > 1:
+            k = k[:, :, None, :, :].expand(B, self.num_kv_heads, self.n_rep, T, self.head_dim)
+            v = v[:, :, None, :, :].expand(B, self.num_kv_heads, self.n_rep, T, self.head_dim)
+            k = k.reshape(B, self.num_heads, T, self.head_dim)
+            v = v.reshape(B, self.num_heads, T, self.head_dim)
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.dropout_prob if self.training else 0.0,
+            is_causal=self.causal
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, C)
+        return attn_output
+
+
+class MultiheadLatentAttentionMixer(Layer):  # Changed Layer to nn.Module for standard torch
+    __name__ = "MultiheadLatentAttentionMixer"
+    # Adjusted complexity notation
+    __complexity__ = "O(L^2 d + L d * d_kv_lora)"
+
+    def __init__(self, d_model: int, n_heads: int, d_kv_lora: int, causal: bool, rope=False, dropout=0.1):
+        super().__init__()
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.d_kv_lora = d_kv_lora
+        self.rope_enabled = rope
+
+        self.q = nn.Linear(d_model, d_model, bias=False)
+
+        self.kv_compress = nn.Linear(d_model, d_kv_lora, bias=False)
+        self.k_content_decompress = nn.Linear(d_kv_lora, d_model, bias=False)
+        self.v_decompress = nn.Linear(d_kv_lora, d_model, bias=False)
+
+        self.out = nn.Linear(d_model, d_model, bias=False)
+
+        if self.rope_enabled:
+            self.rope = RoPE(self.head_dim)
+
+        self.dropout = nn.Dropout(dropout)
+
+        nn.init.zeros_(self.out.weight)
+
+    def forward(self, x: torch.Tensor):
+        B, L, D = x.shape
+        H = self.n_heads
+        HD = self.head_dim
+
+        q = self.q(x).view(B, L, H, HD).transpose(1, 2)
+
+        c_kv = self.kv_compress(x)
+        k = self.k_content_decompress(c_kv).view(B, L, H, HD).transpose(1, 2)
+        v = self.v_decompress(c_kv).view(B, L, H, HD).transpose(1, 2)
+        if self.rope_enabled:
+            q, k = self.rope(q, k)
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.dropout.p if self.training else 0.0,
+            is_causal=self.causal
+        )
+
+        attn = attn.transpose(1, 2).contiguous().view(B, L, D)
+
+        return self.out(attn)
+
+
+class CausalMultiheadAttentionMixer2NOK(nn.Module):
+    __name__ = "CausalMultiheadAttentionMixer"
+
+    def __init__(self, d_model: int, n_heads: int, causal=True, rope=False, dropout=0.0):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+
+        self.qv = nn.Linear(d_model, 2 * d_model, bias=False)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+
+        self.rope = RoPE(self.head_dim) if rope else None
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.qv.weight)
+        nn.init.xavier_uniform_(self.out.weight)
+
+    def forward(self, x: torch.Tensor):
+        B, L, D = x.shape
+
+        qv = self.qv(x)
+
+        qv = qv.view(B, L, 2, self.n_heads, self.head_dim)
+
+        q, v = qv.unbind(dim=2)
+
+        q = q.transpose(1, 2)
+        k = x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        if self.rope is not None:
+            q, k = self.rope(q, k)
+
+        q = q.contiguous()
+        k = k.contiguous()
+        v = v.contiguous()
+
+        attn = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=True,
+            dropout_p=self.dropout_p if self.training else 0.0
+        )
+
+        attn = attn.transpose(1, 2).contiguous().view(B, L, D)
+
+        return self.out(attn)
+
+
+class MultiheadDecoupledAttention(Layer):
+    __name__ = "DecoupledSelfAttention"
+    __complexity__ = "O(L^2 * d_qk + L * d_model * (d_qk + d_v))"
+
+    def __init__(self, d_model: int, n_heads: int, qk_dim: int, causal: bool, rope=False, dropout=0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.qk_dim = qk_dim
+        self.v_dim = d_model // n_heads
+        self.causal = causal
+        self.rope_enabled = rope
+
+        self.qk_proj = nn.Linear(d_model, n_heads * 2 * qk_dim, bias=False)
+
+        self.v_proj = nn.Linear(d_model, n_heads * self.v_dim, bias=False)
+
+        self.out_proj = nn.Linear(n_heads * self.v_dim, d_model, bias=False)
+
+        self.dropout = nn.Dropout(dropout)
+
+        if self.rope_enabled:
+            self.rope = RoPE(qk_dim)
+
+        nn.init.zeros_(self.out_proj.weight)
+
+    def forward(self, x: torch.Tensor):
+        B, L, _ = x.shape
+        H = self.n_heads
+
+        q, k = self.qk_proj(x).view(B, L, H, 2 * self.qk_dim).transpose(1, 2).chunk(2, -1)
+
+        v = self.v_proj(x).view(B, L, H, self.v_dim).transpose(1, 2)
+
+        if self.rope_enabled:
+            q, k = self.rope(q, k)
+
+        attn_out = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.dropout.p if self.training else 0.0,
+            is_causal=self.causal
+        )
+
+        attn_out = attn_out.transpose(1, 2).contiguous().view(B, L, H * self.v_dim)
+
+        return self.out_proj(attn_out)
+
+
+class MultiheadDecoupledAttentionNOV(Layer):
+    __name__ = "DecoupledSelfAttention"
+    __complexity__ = "O(L^2 * d_qk + L * d_model * (d_qk + d_v))"
+
+    def __init__(self, d_model: int, n_heads: int, qk_dim: int, causal: bool, rope=False, dropout=0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.qk_dim = qk_dim
+        self.v_dim = d_model // n_heads
+        self.causal = causal
+        self.rope_enabled = rope
+
+        self.qk_proj = nn.Linear(d_model, n_heads * 2 * qk_dim, bias=False)
+
+        self.out_proj = nn.Linear(n_heads * self.v_dim, d_model, bias=False)
+
+        self.dropout = nn.Dropout(dropout)
+
+        if self.rope_enabled:
+            self.rope = RoPE(qk_dim)
+
+        nn.init.zeros_(self.out_proj.weight)
+
+    def forward(self, x: torch.Tensor):
+        B, L, _ = x.shape
+        H = self.n_heads
+
+        q, k = self.qk_proj(x).view(B, L, H, 2 * self.qk_dim).transpose(1, 2).chunk(2, -1)
+
+        v = x.view(B, L, H, self.v_dim).transpose(1, 2)
+
+        if self.rope_enabled:
+            q, k = self.rope(q, k)
+
+        attn_out = F.scaled_dot_product_attention(
+            q, k, v,
+            dropout_p=self.dropout.p if self.training else 0.0,
+            is_causal=self.causal
+        )
+
+        attn_out = attn_out.transpose(1, 2).contiguous().view(B, L, H * self.v_dim)
+
+        return self.out_proj(attn_out)
+
+
+class FFMultiheadAttentionMixerNOV(Layer):
+    __name__ = "MultiheadAttentionMixer"
+    __complexity__ = "O(L^2 d +  L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool, rope=False, dropout=0.05):
+        super().__init__()
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+
+        self.q = nn.Linear(d_model, d_model, bias=False)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+        if rope:
+            self.rope = RoPE(self.head_dim)
+        else:
+            self.rope = None
+
+        self.dropout = nn.Dropout(dropout)
+
+        self.kv_cache = None
+
+    def forward(self, x: torch.Tensor):
+        B, L, D = x.shape
+
+        q = self.q(x).reshape(B, L, self.n_heads, self.head_dim)
+        q = q.permute(0, 2, 1, 3)
+        k = x.view(B, L, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
+        v = x.view(B, L, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
+
+        if self.rope is not None:
+            q, k = self.rope(q, k, start_pos=0)
+
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal)
+        attn = attn.transpose(1, 2).reshape(B, L, D)
+        attn = self.dropout(attn)
+        return self.out(attn)
+
+    def generate(self, x: torch.Tensor):
+        B, L, D = x.shape
+
+        qk = self.qk(x).reshape(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.permute(2, 0, 3, 1, 4)
+        v = x.view(B, L, 1, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)[0]
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            pos_offset = k_cache.shape[2]
+        else:
+            pos_offset = 0
+            k_cache, v_cache = None, None
+
+        if self.rope is not None:
+            q, k = self.rope(q, k, start_pos=pos_offset)
+
+        if k_cache is not None:
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+        use_causal = self.causal and (L > 1)
+
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=use_causal)
+
+        attn = attn.transpose(1, 2).reshape(B, L, D)
+        attn = self.dropout(attn)
+        return self.out(attn)
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+    def finetune(self):
+        self.qk.weight.data.requires_grad = False
+        self.out.weight.data.requires_grad = False
+
+
+class FFFMultiheadAttentionMixer(Layer):
+    __name__ = "MultiheadAttentionMixer"
+    __complexity__ = "O(L^2 d +  L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool, rope=False, dropout=0.05):
+        super().__init__()
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+
+        self.q = LowRankLinear(d_model, d_model, rank=math.isqrt(d_model), bias=False)
+        self.out = nn.Linear(d_model, d_model, bias=False)
+        if rope:
+            self.rope = RoPE(self.head_dim)
+        else:
+            self.rope = None
+
+        self.dropout = nn.Dropout(dropout)
+
+        self.kv_cache = None
+
+    def forward(self, x: torch.Tensor):
+        B, L, D = x.shape
+
+        q = self.q(x).reshape(B, L, self.n_heads, self.head_dim)
+        q = q.permute(0, 2, 1, 3)
+        k = x.view(B, L, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
+        v = x.view(B, L, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
+
+        if self.rope is not None:
+            q, k = self.rope(q, k, start_pos=0)
+
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal)
+        attn = attn.transpose(1, 2).reshape(B, L, D)
+        attn = self.dropout(attn)
+        return self.out(attn)
+
+    def generate(self, x: torch.Tensor):
+        B, L, D = x.shape
+
+        qk = self.qk(x).reshape(B, L, 2, self.n_heads, self.head_dim)
+        q, k = qk.permute(2, 0, 3, 1, 4)
+        v = x.view(B, L, 1, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)[0]
+
+        if self.kv_cache is not None:
+            k_cache, v_cache = self.kv_cache
+            pos_offset = k_cache.shape[2]
+        else:
+            pos_offset = 0
+            k_cache, v_cache = None, None
+
+        if self.rope is not None:
+            q, k = self.rope(q, k, start_pos=pos_offset)
+
+        if k_cache is not None:
+            k = torch.cat([k_cache, k], dim=2)
+            v = torch.cat([v_cache, v], dim=2)
+
+        self.kv_cache = (k, v)
+        use_causal = self.causal and (L > 1)
+
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=use_causal)
+
+        attn = attn.transpose(1, 2).reshape(B, L, D)
+        attn = self.dropout(attn)
+        return self.out(attn)
+
+    def clear_cache(self):
+        self.kv_cache = None
+
+    def finetune(self):
+        self.qk.weight.data.requires_grad = False
+        self.out.weight.data.requires_grad = False
+
+
+class NSAttention(Layer):
+    __name__ = "MultiheadAttentionMixer"
+    __complexity__ = "O(L^2 d + L d^2)"
+
+    def __init__(self, d_model: int, n_heads: int, causal: bool = True,
+                 pos_emb: nn.Module = None, dropout: float = 0.0, qk_norm=False):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.causal = causal
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.dropout_p = dropout
+
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+
+        self.pos_emb = pos_emb
+        self.kv_cache = None
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.qkv.weight)
+        nn.init.xavier_uniform_(self.out.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, L, D = x.shape
+
+        qkv = self.qkv(x)
+
+        q, k, v = qkv.chunk(3, dim=-1)
+        kv = k.transpose(-1, -2) @ v
+        return q @ kv.transpose(-1, -2)
+
+
+# MODDED NANO GPT
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, rope_base: float, qk_gain_init: float):
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError("model_dim must be divisible by num_heads")
+
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        if self.head_dim % 2 != 0:
+            raise ValueError("head_dim must be even for RoPE")
+
+        # FUSED QKV: Output dimension is 3 * dim
+        self.c_qkv = CastedLinear(dim, 3 * dim, bias=False)
+
+        self.proj = CastedLinear(dim, dim, bias=False)
+        self.proj._zero_init = True
+
+        self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32))
+        self.rotary = Rotary(self.head_dim, base=rope_base)
+
+    def forward(self, x: Tensor) -> Tensor:
+        bsz, seqlen, dim = x.shape
+
+        # 1. Single fused linear projection
+        # Shape: (bsz, seqlen, 3 * dim)
+        qkv = self.c_qkv(x)
+
+        # 2. Reshape and slice into Q, K, V
+        # Reshape to: (bsz, seqlen, 3, num_heads, head_dim)
+        qkv = qkv.reshape(bsz, seqlen, 3, self.num_heads, self.head_dim)
+        # Permute to: (3, bsz, num_heads, seqlen, head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+
+        # Unbind separates the first dimension (the '3') into a tuple
+        q, k, v = qkv.unbind(0)
+
+        # 3. Apply norms
+        q = F.rms_norm(q, (q.size(-1),))
+        k = F.rms_norm(k, (k.size(-1),))
+
+        # 4. RoPE
+        cos, sin = self.rotary(seqlen, x.device, q.dtype)
+        q = apply_rotary_emb(q, cos, sin)
+        k = apply_rotary_emb(k, cos, sin)
+
+        # 5. Q-gain scaling
+        q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None]
+
+        # 6. Attention
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None,
+            is_causal=True,
+            enable_gqa=False,  # Standard MHA
+        )
+
+        # 7. Final projection
+        y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim)
+        return self.proj(y)

vathos/_triton.py

@@ -0,0 +1,530 @@
+"""
+Triton-fused two-layer MLPs in the style of `VariableUDLP`:
+    out = (act(x @ W1.T))^2 @ W2.T          (ReLU^2, LeakyReLU(0.5)^2)
+    out = (silu(a) * b) @ W2.T               (SwiGLU)
+
+The fused ReLU^2 and LeakyReLU^2 kernels are ports of:
+  - modded-nanogpt   linear_relu_square        (Keller Jordan et al.)
+  - parameter-golf   linear_leaky_relu_square  (samacqua / openai/parameter-golf)
+
+Both use the Hopper TensorDescriptor (TMA) API, so they need:
+  - triton >= 3.2
+  - GPU compute capability >= 9.0 (H100 / GH200)
+On older GPUs / when triton is missing, the wrapper modules transparently fall
+back to a pure-torch path that is `torch.compile`-friendly.
+
+The autograd.Function wrappers are accepted natively by `torch.compile` (>=2.0):
+the compiler treats `.apply(...)` as an opaque op and traces around it.
+
+API: each module mirrors `VariableUDLP(d_model, d_output, M, dropout)`.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from Vathos._basics import Layer
+from Vathos.functions import flag
+
+try:
+    import triton
+    import triton.language as tl
+    from triton.tools.tensor_descriptor import TensorDescriptor
+    _HAS_TRITON = True
+except Exception:
+    _HAS_TRITON = False
+
+
+def _hopper_available() -> bool:
+    if not (_HAS_TRITON and torch.cuda.is_available()):
+        return False
+    major, _ = torch.cuda.get_device_capability()
+    return major >= 9
+
+
+# =============================================================================
+# Triton kernels
+# =============================================================================
+
+if _HAS_TRITON:
+
+    @triton.jit
+    def _linear_relu_square_kernel(
+        a_desc, b_desc, c_desc, aux_desc,
+        M, N, K,
+        BLOCK_SIZE_M: tl.constexpr,
+        BLOCK_SIZE_N: tl.constexpr,
+        BLOCK_SIZE_K: tl.constexpr,
+        NUM_SMS: tl.constexpr,
+        FORWARD: tl.constexpr,
+    ):
+        """Port of modded-nanogpt's linear_relu_square_kernel.
+
+        FORWARD:  computes c = (x @ W1.T) and aux = relu(c)^2  (c is the pre-activation)
+        BACKWARD: computes c = grad_pre = grad_post * 2 * relu(pre), reading pre from aux
+                  (caller passes aux=pre, grad_post is the input matmul result)
+        """
+        dtype = tl.bfloat16
+        start_pid = tl.program_id(axis=0)
+        num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+        num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+        num_tiles = num_pid_m * num_pid_n
+        k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
+
+        for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True):
+            pid_m = tile_id // num_pid_n
+            pid_n = tile_id % num_pid_n
+            offs_am = pid_m * BLOCK_SIZE_M
+            offs_bn = pid_n * BLOCK_SIZE_N
+
+            accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+            for ki in range(k_tiles):
+                offs_k = ki * BLOCK_SIZE_K
+                a = a_desc.load([offs_am, offs_k])
+                b = b_desc.load([offs_bn, offs_k])
+                accumulator = tl.dot(a, b.T, accumulator)
+
+            acc = tl.reshape(accumulator, (BLOCK_SIZE_M, 2, BLOCK_SIZE_N // 2))
+            acc = tl.permute(acc, (0, 2, 1))
+            acc0, acc1 = tl.split(acc)
+
+            c0 = acc0.to(dtype)
+            if not FORWARD:
+                pre0 = aux_desc.load([offs_am, offs_bn])
+                c0 = 2 * c0 * tl.where(pre0 > 0, pre0, 0)
+            c_desc.store([offs_am, offs_bn], c0)
+            if FORWARD:
+                p0 = tl.maximum(c0, 0)
+                aux_desc.store([offs_am, offs_bn], p0 * p0)
+
+            c1 = acc1.to(dtype)
+            if not FORWARD:
+                pre1 = aux_desc.load([offs_am, offs_bn + BLOCK_SIZE_N // 2])
+                c1 = 2 * c1 * tl.where(pre1 > 0, pre1, 0)
+            c_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], c1)
+            if FORWARD:
+                p1 = tl.maximum(c1, 0)
+                aux_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], p1 * p1)
+
+
+    @triton.jit
+    def _linear_leaky_relu_square_kernel(
+        a_desc, b_desc, c_desc, aux_desc,
+        M, N, K,
+        BLOCK_SIZE_M: tl.constexpr,
+        BLOCK_SIZE_N: tl.constexpr,
+        BLOCK_SIZE_K: tl.constexpr,
+        NUM_SMS: tl.constexpr,
+        FORWARD: tl.constexpr,
+    ):
+        """Port of parameter-golf's linear_leaky_relu_square_kernel.
+        Uses LeakyReLU with negative slope 0.5, then squares.
+        Forward: aux = lrelu(c)^2,    Backward: c = grad_post * 2 * (lrelu'(pre) * lrelu(pre))
+                                                = grad_post * (where(pre>0, 2*pre, 0.5*pre))
+        """
+        dtype = tl.bfloat16
+        start_pid = tl.program_id(axis=0)
+        num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+        num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+        num_tiles = num_pid_m * num_pid_n
+        k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
+
+        for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True):
+            pid_m = tile_id // num_pid_n
+            pid_n = tile_id % num_pid_n
+            offs_am = pid_m * BLOCK_SIZE_M
+            offs_bn = pid_n * BLOCK_SIZE_N
+
+            accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+            for ki in range(k_tiles):
+                offs_k = ki * BLOCK_SIZE_K
+                a = a_desc.load([offs_am, offs_k])
+                b = b_desc.load([offs_bn, offs_k])
+                accumulator = tl.dot(a, b.T, accumulator)
+
+            acc = tl.reshape(accumulator, (BLOCK_SIZE_M, 2, BLOCK_SIZE_N // 2))
+            acc = tl.permute(acc, (0, 2, 1))
+            acc0, acc1 = tl.split(acc)
+
+            c0 = acc0.to(dtype)
+            c1 = acc1.to(dtype)
+            if not FORWARD:
+                pre0 = aux_desc.load([offs_am, offs_bn])
+                pre1 = aux_desc.load([offs_am, offs_bn + BLOCK_SIZE_N // 2])
+                c0 = c0 * tl.where(pre0 > 0, 2.0 * pre0, 0.5 * pre0)
+                c1 = c1 * tl.where(pre1 > 0, 2.0 * pre1, 0.5 * pre1)
+            c_desc.store([offs_am, offs_bn], c0)
+            c_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], c1)
+            if FORWARD:
+                aux0 = tl.where(c0 > 0, c0, 0.5 * c0)
+                aux1 = tl.where(c1 > 0, c1, 0.5 * c1)
+                aux_desc.store([offs_am, offs_bn], aux0 * aux0)
+                aux_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], aux1 * aux1)
+
+
+# -----------------------------------------------------------------------------
+# T4-friendly LeakyReLU(0.5)^2 kernel (no TMA, fp16, smaller tiles).
+# Designed for sm_75 (Turing, T4): pointer-based loads, BLOCK_M=BLOCK_N=64,
+# BLOCK_K=32, fp16 inputs/outputs, fp32 accumulator. Same fwd/bwd convention
+# as the Hopper kernels: forward stores c=pre and aux=post; backward reads
+# pre from aux and writes grad_pre into c. Output dtype follows the inputs
+# (cast to fp32 fails tl.dot on Turing — keep them fp16).
+# -----------------------------------------------------------------------------
+
+if _HAS_TRITON:
+
+    @triton.jit
+    def _linear_lrelu2_kernel_t4(
+        a_ptr, b_ptr, c_ptr, aux_ptr,
+        M, N, K,
+        stride_am, stride_ak,
+        stride_bn, stride_bk,
+        stride_cm, stride_cn,
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        BLOCK_K: tl.constexpr,
+        GROUP_M: tl.constexpr,
+        FORWARD: tl.constexpr,
+    ):
+        pid = tl.program_id(axis=0)
+        num_pid_m = tl.cdiv(M, BLOCK_M)
+        num_pid_n = tl.cdiv(N, BLOCK_N)
+        num_pid_in_group = GROUP_M * num_pid_n
+        group_id = pid // num_pid_in_group
+        first_pid_m = group_id * GROUP_M
+        group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_M)
+        pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+        pid_n = (pid % num_pid_in_group) // group_size_m
+
+        offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        offs_k = tl.arange(0, BLOCK_K)
+
+        a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
+        b_ptrs = b_ptr + offs_n[:, None] * stride_bn + offs_k[None, :] * stride_bk
+
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k in range(0, tl.cdiv(K, BLOCK_K)):
+            k_remaining = K - k * BLOCK_K
+            a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < k_remaining)
+            b_mask = (offs_n[:, None] < N) & (offs_k[None, :] < k_remaining)
+            a = tl.load(a_ptrs, mask=a_mask, other=0.0)
+            b = tl.load(b_ptrs, mask=b_mask, other=0.0)
+            acc += tl.dot(a, b.T)
+            a_ptrs += BLOCK_K * stride_ak
+            b_ptrs += BLOCK_K * stride_bk
+
+        out_dtype = c_ptr.dtype.element_ty
+        c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
+        aux_ptrs = aux_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
+        mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+
+        if FORWARD:
+            # acc here is pre = x @ W1.T. Save pre into c, post = lrelu²(pre) into aux.
+            tl.store(c_ptrs, acc.to(out_dtype), mask=mask)
+            scale = tl.where(acc > 0, 1.0, 0.25)   # act² = z² * (1 if z>0 else 0.25)
+            post = acc * acc * scale
+            tl.store(aux_ptrs, post.to(out_dtype), mask=mask)
+        else:
+            # acc here is grad_post @ W2 (b=W2.T). Multiply elementwise by
+            # d(lrelu²)/dpre = where(pre>0, 2*pre, 0.5*pre) using saved pre in aux.
+            pre = tl.load(aux_ptrs, mask=mask, other=0.0).to(tl.float32)
+            deriv = tl.where(pre > 0, 2.0 * pre, 0.5 * pre)
+            tl.store(c_ptrs, (acc * deriv).to(out_dtype), mask=mask)
+
+
+def _launch_lrelu2_t4(a: torch.Tensor, b: torch.Tensor, aux: torch.Tensor | None):
+    """Pointer-based kernel launcher for T4. Returns (pre, post) in fwd, grad_pre in bwd."""
+    M, K = a.shape
+    N, K2 = b.shape
+    assert K == K2
+
+    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
+    forward = aux is None
+    if aux is None:
+        aux = torch.empty((M, N), device=a.device, dtype=a.dtype)
+
+    BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
+    grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
+    _linear_lrelu2_kernel_t4[grid](
+        a, b, c, aux,
+        M, N, K,
+        a.stride(0), a.stride(1),
+        b.stride(0), b.stride(1),
+        c.stride(0), c.stride(1),
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        GROUP_M=8,
+        FORWARD=forward,
+        num_warps=4,
+        num_stages=2,
+    )
+    return (c, aux) if forward else c
+
+
+class _FusedLinearLReLU2_T4(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, W1, W2):
+        x_flat = x.reshape(-1, x.shape[-1])
+        pre, post = _launch_lrelu2_t4(x_flat, W1, None)
+        out = post @ W2.T
+        ctx.save_for_backward(x, W1, W2, pre, post)
+        return out.view(*x.shape[:-1], W2.shape[0])
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        x, W1, W2, pre, post = ctx.saved_tensors
+        x_flat = x.reshape(-1, x.shape[-1])
+        go = grad_out.reshape(-1, W2.shape[0])
+        dW2 = go.T @ post
+        dpre = _launch_lrelu2_t4(go, W2.T.contiguous(), pre)
+        dW1 = dpre.T @ x_flat
+        dx = dpre @ W1
+        return dx.view_as(x), dW1, dW2
+
+
+def _can_use_triton_t4(x) -> bool:
+    """T4 (CC 7.5) requires fp16 inputs."""
+    if not (_HAS_TRITON and torch.cuda.is_available()):
+        return False
+    major, minor = torch.cuda.get_device_capability()
+    if (major, minor) != (7, 5):
+        return False
+    return x.dtype == torch.float16
+
+
+# =============================================================================
+# Host-side launchers (return both pre-activation and post-activation in fwd)
+# =============================================================================
+
+def _launch_fused_act_square(kernel_fn, a: torch.Tensor, b: torch.Tensor,
+                             aux: torch.Tensor | None):
+    """Shared launcher for both relu^2 and leaky_relu^2 kernels.
+    Forward (aux=None): returns (pre, post) = (x@W.T, act(pre)^2)
+    Backward (aux=pre): returns dpre = grad_post * d(act(pre)^2)/dpre
+    """
+    M, K = a.shape
+    N, K2 = b.shape
+    assert K == K2
+
+    c = torch.empty((M, N), device=a.device, dtype=a.dtype)
+    forward = aux is None
+    if aux is None:
+        aux = torch.empty((M, N), device=a.device, dtype=a.dtype)
+
+    BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 256, 64
+    num_stages = 4 if forward else 3
+
+    a_desc = TensorDescriptor.from_tensor(a, [BLOCK_SIZE_M, BLOCK_SIZE_K])
+    b_desc = TensorDescriptor.from_tensor(b, [BLOCK_SIZE_N, BLOCK_SIZE_K])
+    c_desc = TensorDescriptor.from_tensor(c, [BLOCK_SIZE_M, BLOCK_SIZE_N // 2])
+    aux_desc = TensorDescriptor.from_tensor(aux, [BLOCK_SIZE_M, BLOCK_SIZE_N // 2])
+
+    num_sms = torch.cuda.get_device_properties(a.device).multi_processor_count
+    grid = lambda _meta: (
+        min(num_sms, triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N)),
+    )
+    kernel_fn[grid](
+        a_desc, b_desc, c_desc, aux_desc,
+        M, N, K,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=BLOCK_SIZE_N,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+        NUM_SMS=num_sms,
+        FORWARD=forward,
+        num_stages=num_stages,
+        num_warps=8,
+    )
+    return (c, aux) if forward else c
+
+
+# =============================================================================
+# Autograd Functions — composes the fused kernel with the second matmul.
+# Compatible with torch.compile (autograd.Function is opaque to dynamo).
+# =============================================================================
+
+class _FusedLinearReLU2(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, W1, W2):
+        x_flat = x.reshape(-1, x.shape[-1])
+        pre, post = _launch_fused_act_square(_linear_relu_square_kernel, x_flat, W1, None)
+        out = post @ W2.T
+        ctx.save_for_backward(x, W1, W2, pre, post)
+        return out.view(*x.shape[:-1], W2.shape[0])
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        x, W1, W2, pre, post = ctx.saved_tensors
+        x_flat = x.reshape(-1, x.shape[-1])
+        go = grad_out.reshape(-1, W2.shape[0])
+        dW2 = go.T @ post
+        dpre = _launch_fused_act_square(_linear_relu_square_kernel, go, W2.T.contiguous(), pre)
+        dW1 = dpre.T @ x_flat
+        dx = dpre @ W1
+        return dx.view_as(x), dW1, dW2
+
+
+class _FusedLinearLReLU2(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, W1, W2):
+        x_flat = x.reshape(-1, x.shape[-1])
+        pre, post = _launch_fused_act_square(_linear_leaky_relu_square_kernel, x_flat, W1, None)
+        out = post @ W2.T
+        ctx.save_for_backward(x, W1, W2, pre, post)
+        return out.view(*x.shape[:-1], W2.shape[0])
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        x, W1, W2, pre, post = ctx.saved_tensors
+        x_flat = x.reshape(-1, x.shape[-1])
+        go = grad_out.reshape(-1, W2.shape[0])
+        dW2 = go.T @ post
+        dpre = _launch_fused_act_square(_linear_leaky_relu_square_kernel, go, W2.T.contiguous(), pre)
+        dW1 = dpre.T @ x_flat
+        dx = dpre @ W1
+        return dx.view_as(x), dW1, dW2
+
+
+# =============================================================================
+# Pure-torch fallbacks (torch.compile-friendly)
+# =============================================================================
+
+def _eager_relu2_mlp(x, W1, W2):
+    return F.relu(F.linear(x, W1)).square() @ W2.T
+
+def _eager_lrelu2_mlp(x, W1, W2):
+    return F.leaky_relu(F.linear(x, W1), 0.5).square() @ W2.T
+
+
+def _can_use_triton_fused(x, M, N) -> bool:
+    """The TMA kernels require shapes divisible by their tile sizes and a Hopper GPU."""
+    if not _hopper_available():
+        return False
+    if x.dtype not in (torch.bfloat16, torch.float16):
+        return False
+    return (M % 128 == 0) and (N % 256 == 0)
+
+
+# =============================================================================
+# VariableUDLP-compatible nn.Modules
+# =============================================================================
+
+class _BaseTritonUDLP(Layer):
+    """Shared structure: expand (d_model→M, no bias) -> fused act² -> contract (M→d_output, no bias).
+
+    Matches `VariableUDLP(d_model, d_output, M, dropout)` so that it can be passed
+    as the `UDLP=` argument of `ModdedFormer` without further changes. The
+    `activation` parameter accepted by `VariableUDLP` is ignored here — the
+    activation is fixed by the chosen kernel (subclass).
+    """
+
+    _autograd_fn = None
+    _eager_fn = None
+    _name = "TritonUDLP"
+
+    def __init__(self, d_model: int, d_output: int, M: int,
+                 activation=None, dropout: float = 0.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_output = d_output
+        self.M = M
+        self.expand = nn.Linear(d_model, M, bias=False)
+        self.contract = nn.Linear(M, d_output, bias=False)
+        self.dropout = nn.Dropout(dropout)
+        if activation is not None:
+            flag(f"{self._name}: 'activation' arg ignored — kernel hardcodes the activation.", 3)
+
+    def _init_weights(self):
+        torch.nn.init.zeros_(self.contract.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        W1 = self.expand.weight       # [M, d_model]
+        W2 = self.contract.weight     # [d_output, M]
+        if _can_use_triton_fused(x, self.M, self.d_output):
+            out = self._autograd_fn.apply(x, W1, W2)
+        else:
+            out = self._eager_fn(x, W1, W2)
+        return self.dropout(out)
+
+
+class TritonReLU2UDLP(_BaseTritonUDLP):
+    """Fused ReLU² UDLP. Drop-in for VariableUDLP with activation=ReLU2."""
+    _autograd_fn = _FusedLinearReLU2 if _HAS_TRITON else None
+    _eager_fn = staticmethod(_eager_relu2_mlp)
+    _name = "TritonReLU2UDLP"
+
+
+class TritonLReLU2UDLP(_BaseTritonUDLP):
+    """Fused LeakyReLU(0.5)² UDLP. Drop-in for VariableUDLP with activation=LeakyReLU2."""
+    _autograd_fn = _FusedLinearLReLU2 if _HAS_TRITON else None
+    _eager_fn = staticmethod(_eager_lrelu2_mlp)
+    _name = "TritonLReLU2UDLP"
+
+
+class TritonLReLU2UDLP_T4(_BaseTritonUDLP):
+    """LeakyReLU(0.5)² UDLP tuned for T4 / Turing (sm_75).
+
+    Uses a non-TMA fp16 kernel with 64×64×32 tiles. Falls back to the eager
+    torch path on non-T4 GPUs or when triton is missing.
+
+    NOTE: weights are kept fp32 by Linear; cast inputs to fp16 before forward
+    (e.g. `model = model.half()` or `with torch.autocast('cuda', dtype=torch.float16)`).
+    """
+    _autograd_fn = _FusedLinearLReLU2_T4 if _HAS_TRITON else None
+    _eager_fn = staticmethod(_eager_lrelu2_mlp)
+    _name = "TritonLReLU2UDLP_T4"
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        W1 = self.expand.weight
+        W2 = self.contract.weight
+        if _can_use_triton_t4(x):
+            out = self._autograd_fn.apply(x, W1, W2)
+        else:
+            out = self._eager_fn(x, W1, W2)
+        return self.dropout(out)
+
+
+# -----------------------------------------------------------------------------
+# SwiGLU — torch-only path (compile-friendly).
+#
+# SwiGLU is gated, so the expand projects to 2*M and the activation consumes
+# both halves: silu(a) * b. There is no clean port of this in modded-nanogpt
+# or parameter-golf; Liger-Kernel implements a full fused fwd+bwd for it but
+# that's a separate dependency. The torch path below compiles cleanly and is
+# usually within ~10% of a custom kernel for the typical M sizes we use.
+# -----------------------------------------------------------------------------
+
+class TritonSwiGLUUDLP(Layer):
+    """SwiGLU UDLP, drop-in for VariableUDLP with activation=SwiGLU.
+    Compile-friendly torch implementation — no custom kernel.
+    `M` is the post-gate inner dimension (i.e. expand projects to 2*M).
+    """
+
+    def __init__(self, d_model: int, d_output: int, M: int,
+                 activation=None, dropout: float = 0.0):
+        super().__init__()
+        self.d_model = d_model
+        self.d_output = d_output
+        self.M = M
+        self.expand = nn.Linear(d_model, 2 * M, bias=False)
+        self.contract = nn.Linear(M, d_output, bias=False)
+        self.dropout = nn.Dropout(dropout)
+        if activation is not None:
+            flag("TritonSwiGLUUDLP: 'activation' arg ignored — gate is fixed to SiLU.", 3)
+
+    def _init_weights(self):
+        torch.nn.init.zeros_(self.contract.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.expand(x)
+        a, b = h.chunk(2, dim=-1)
+        return self.dropout(self.contract(F.silu(a) * b))
+
+
+__all__ = [
+    "TritonReLU2UDLP",
+    "TritonLReLU2UDLP",
+    "TritonLReLU2UDLP_T4",
+    "TritonSwiGLUUDLP",
+]

vathos/blocks.py

@@ -0,0 +1,2110 @@
+"""
+All the layers included here aim to be already optimized by torch itself, without requiring triton/cuda kernels (explicitly),
+in fact they try to only use matmul and Linear operations which are arguably already optimized implicitly in torch.
+"""
+import time
+
+import torch
+
+from Vathos._basics import *
+from typing import Tuple, Optional, Union, List
+import re
+from Vathos.complexity import combine_big_o_product, combine_big_o_sum
+import math
+import torch.nn.functional as F
+from tqdm import tqdm
+from Vathos._spatials import *
+
+
+class Block1d(Layer):
+    def __init__(self, d_model, channel_mixer: Layer, spatial_mixer: Layer, norm=nn.LayerNorm):
+        super().__init__()
+        self.spatial_mixer = spatial_mixer
+        self.channel_mixer = channel_mixer
+        self.norm1 = norm(spatial_mixer.d_model)
+        self.norm2 = norm(spatial_mixer.d_model)
+
+    def forward(self, x: torch.Tensor, ve=None):
+        if ve is not None:
+            x = x + self.spatial_mixer(self.norm1(x), ve)
+        else:
+            x = x + self.spatial_mixer(self.norm1(x))
+
+        x = x + self.channel_mixer(self.norm2(x))
+        return x
+
+    def generate(self, x: torch.Tensor, ve=None):
+        h = self.norm1(x)
+        if self.spatial_mixer.has_custom_generate():
+            spatial_out = _call_with_ve(self.spatial_mixer.generate, h, ve)
+        else:
+            spatial_out = _call_with_ve(self.spatial_mixer, h, ve)
+        x = x + spatial_out
+
+        if self.channel_mixer.has_custom_generate():
+            x = x + self.channel_mixer.generate(self.norm2(x))
+        else:
+            x = x + self.channel_mixer(self.norm2(x))
+
+        return x
+
+
+def _call_with_ve(fn, h, ve):
+    """Call fn(h, ve=ve) if it accepts ve, otherwise fn(h). Used by Block.generate
+    so that spatial mixers without ve support stay drop-in."""
+    if ve is None:
+        return fn(h)
+    try:
+        return fn(h, ve=ve)
+    except TypeError:
+        return fn(h)
+
+
+class CBlock1d(Layer):
+    def __init__(self, d_model, channel_mixer: Layer, spatial_mixer: Layer, norm=nn.LayerNorm):
+        super().__init__()
+        self.spatial_mixer = spatial_mixer
+        self.channel_mixer = channel_mixer
+        self.norm1 = norm(spatial_mixer.d_model)
+        self.norm2 = norm(spatial_mixer.d_model)
+        self.l = nn.Parameter(torch.tensor([1.0, 1.0, 1.0, 1.0]))
+
+    def forward(self, x: torch.Tensor, ve=None):
+        if ve is not None:
+            x = x * self.l[0] + self.spatial_mixer(self.norm1(x), ve) * self.l[1]
+        else:
+            x = x * self.l[0] + self.spatial_mixer(self.norm1(x)) * self.l[1]
+        x = x * self.l[2] + self.channel_mixer(self.norm2(x)) * self.l[3]
+        return x
+
+    def generate(self, x: torch.Tensor, ve=None):
+        h = self.norm1(x)
+        if self.spatial_mixer.has_custom_generate():
+            spatial_out = _call_with_ve(self.spatial_mixer.generate, h, ve)
+        else:
+            spatial_out = _call_with_ve(self.spatial_mixer, h, ve)
+        x = x * self.l[0] + spatial_out * self.l[1]
+        if self.channel_mixer.has_custom_generate():
+            x = x * self.l[2] + self.channel_mixer.generate(self.norm2(x)) * self.l[3]
+        else:
+            x = x * self.l[2] + self.channel_mixer(self.norm2(x)) * self.l[3]
+        return x
+
+
+# TODO: Smear
+class SmearBlock1d(Layer):
+    def __init__(self, d_model, channel_mixer: Layer, spatial_mixer: Layer, norm=nn.LayerNorm):
+        super().__init__()
+        self.spatial_mixer = spatial_mixer
+        self.channel_mixer = channel_mixer
+        self.norm1 = norm(spatial_mixer.d_model)
+        self.norm2 = norm(spatial_mixer.d_model)
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.spatial_mixer(self.norm1(x))
+        x = x + self.channel_mixer(self.norm2(x))
+        return x
+
+    def generate(self, x: torch.Tensor):
+        # Use generate if available, otherwise forward
+        if self.spatial_mixer.has_custom_generate():
+            x = x + self.spatial_mixer.generate(self.norm1(x))
+        else:
+            x = x + self.spatial_mixer(self.norm1(x))
+
+        if self.channel_mixer.has_custom_generate():
+            x = x + self.channel_mixer.generate(self.norm2(x))
+        else:
+            x = x + self.channel_mixer(self.norm2(x))
+
+        return x
+
+
+class JBlock1d(Layer):
+    def __init__(self, d_model, channel_mixer: nn.Module, spatial_mixer: nn.Module, norm=nn.LayerNorm):
+        super().__init__()
+        self.channel_mixer = channel_mixer
+        self.spatial_mixer = spatial_mixer
+        self.norm = norm(d_model)
+
+    def forward(self, x: torch.Tensor, ve=None):
+        h = self.norm(x)
+        if ve is not None:
+            return x + self.spatial_mixer(h, ve) + self.channel_mixer(h)
+        return x + self.spatial_mixer(h) + self.channel_mixer(h)
+
+    def generate(self, x: torch.Tensor, ve=None):
+        h = self.norm(x)
+        sm_fn = self.spatial_mixer.generate if (hasattr(self.spatial_mixer, 'has_custom_generate') and self.spatial_mixer.has_custom_generate()) else self.spatial_mixer
+        cm_fn = self.channel_mixer.generate if (hasattr(self.channel_mixer, 'has_custom_generate') and self.channel_mixer.has_custom_generate()) else self.channel_mixer
+        spatial_out = _call_with_ve(sm_fn, h, ve)
+        channel_out = cm_fn(h)
+        return x + spatial_out + channel_out
+
+
+class ExpandingBlock1d(Layer):
+    def __init__(self, d_model, channel_mixer: Layer, spatial_mixer: Layer, expand=1):
+        super().__init__()
+        self.spatial_mixer = spatial_mixer
+        self.channel_mixer = channel_mixer
+        self.expander = nn.Linear(d_model, d_model * expand, bias=False)
+        self.contractor = nn.Linear(d_model * expand, d_model, bias=False)
+        self.norm1 = nn.LayerNorm(spatial_mixer.d_model)
+        self.norm2 = nn.LayerNorm(spatial_mixer.d_model)
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.contractor(self.spatial_mixer(self.norm1(self.expander(x))))
+        x = x + self.channel_mixer(self.norm2(x))
+        return x
+
+
+class Renamer:
+    def __init__(self, constructor, renames: dict):
+        super().__init__()
+        self.constructor = constructor
+        self.renames = renames
+
+    def __call__(self, *args, **kwargs):
+        renamed_kwargs = {}
+        for key, value in kwargs.items():
+            if key in self.renames:
+                renamed_kwargs[self.renames[key]] = value
+            else:
+                renamed_kwargs[key] = value
+        return self.constructor(*args, renamed_kwargs)
+
+
+class BlockStack(Layer):
+    def __init__(self, blocks: Tuple[Block1d | Layer]):
+        super().__init__()
+        self.blocks = blocks
+        self.stack = nn.ModuleList(blocks)
+
+    def forward(self, x: torch):
+        for block in self.stack:
+            x = block(x)
+        return x
+
+
+class DepthwiseCausalConv1d(Layer):
+    __name__ = "CausalConv1d"
+    __complexity__ = "O(L d k)"
+
+    def __init__(self, d_model, k=3):
+        super().__init__()
+        self.d_model = d_model
+        self.k = k
+        self.pad = k - 1
+
+        self.K = nn.Parameter(torch.randn(d_model, k) / (k ** 0.5))
+
+    def forward(self, x):
+        b, L, d = x.shape
+        x = x.transpose(1, 2)
+        x_pad = F.pad(x, (self.pad, 0))
+        out = F.conv1d(x_pad, self.K.unsqueeze(1), groups=d)
+        return out.transpose(1, 2)
+
+
+class CausalConv1d(Layer):
+    __name__ = "CausalConv1d"
+    __complexity__ = "O(L d k)"
+
+    def __init__(self, d_model, k=3, groups=None, outproj=False):
+        super().__init__()
+        self.d_model = d_model
+        self.k = k
+        self.pad = k - 1
+        self.groups = groups if groups is not None else d_model
+        self.outproj = nn.Linear(groups, d_model, bias=False) if outproj else Identity()
+        if self.groups != d_model and not outproj:
+            flag(
+                "Output channel numbers will be the number of groups, use outproj=True if you want to enable a linear projection to go back to d_model")
+
+        self.K = nn.Parameter(torch.randn(d_model, k) / (k ** 0.5))
+
+    def forward(self, x):
+        x = x.transpose(1, 2)
+        x_pad = F.pad(x, (self.pad, 0))
+        out = F.conv1d(x_pad, self.K.unsqueeze(1), groups=self.groups)
+        return out.transpose(1, 2)
+
+
+class LSTM(Layer):
+    __name__ = "LSTM"
+    __complexity__ = "O(L d^2)"
+
+    def __init__(self, d_model, d_hidden, bidirectional=False, dropout=0.1, n_layer=1):
+        super().__init__()
+        self.bidirectional = bidirectional
+        self.d_model = d_model
+        self.d_hidden = d_hidden
+        self.projout = nn.Linear(d_hidden, d_model) if d_model != d_hidden else Identity()
+        self.LSTM = nn.LSTM(d_model, d_hidden, bidirectional=bidirectional, batch_first=True, num_layers=n_layer)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        out, _ = self.LSTM(x)
+        out = self.dropout(out)
+        return self.projout(out)
+
+
+class GRU(Layer):
+    __name__ = "GRU"
+    __complexity__ = "O(L d^2)"
+
+    def __init__(self, d_model, d_hidden, bidirectional=False, dropout=0.1, n_layer=1):
+        super().__init__()
+        self.bidirectional = bidirectional
+        self.d_model = d_model
+        self.d_hidden = d_hidden
+        self.projout = nn.Linear(d_hidden, d_model) if d_model != d_hidden else Identity()
+        self.GRU = nn.GRU(d_model, d_hidden, bidirectional=bidirectional, batch_first=True, num_layers=n_layer)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        out, _ = self.GRU(x)
+        out = self.dropout(out)
+        return self.projout(out)
+
+
+class LinearMixer(Layer):
+    __name__ = "Linear Mixer"
+    __complexity__ = "O(L^2 d)"
+
+    def __init__(self, d_model, max_len=1000, causal=True):
+        super().__init__()
+        self.causal = causal
+        self.d_model = d_model
+        self.W = nn.Parameter(torch.randn(max_len, max_len) * 0.01)
+
+    def forward(self, x):
+        x = x / math.sqrt(self.d_model)
+        if self.causal:
+            return torch.tril(self.W[:x.shape[1], :x.shape[1]]) @ x
+        else:
+            return self.W[:x.shape[1], :x.shape[1]] @ x
+
+
+class MLPMixer(Layer):
+    __name__ = "Linear Mixer"
+    __complexity__ = "O(L^2 d)"
+
+    def __init__(self, d_model, max_len=1000, L_expand=1, causal=False, activation=nn.GELU):
+        super().__init__()
+        if causal:
+            if L_expand < 1:
+                raise ValueError("For the MLPMixer to be causal, L_expand must be greater than 1")
+            if L_expand > 1:
+                raise NotImplementedError(
+                    "Causal MLP Mixer with L_expand is currently not implemented")  # TODO: Causal Lex
+        self.causal = causal
+        self.d_model = d_model
+        self.W1 = nn.Parameter(torch.randn(max_len * L_expand, max_len) / max_len)
+        self.W2 = nn.Parameter(torch.randn(max_len, max_len * L_expand) / max_len)
+        self.activation = activation()
+
+    def forward(self, x):
+        x = x / math.sqrt(self.d_model)
+
+        return self.W2 @ self.activation(self.W1 @ x)
+
+
+class ShortConvGatedMixer(Layer):
+    def __init__(self, d_model, mixer, mixer_params, k=4, activation=nn.Sigmoid, k1=None, k2=None):
+        super().__init__()
+        if k1 is None:
+            k1 = k2 = k
+
+        self.mixer = mixer(mixer_params)
+        self.conv_1 = DepthwiseCausalConv1d(d_model, k=k1)
+        self.conv_2 = DepthwiseCausalConv1d(d_model, k=k2)
+        self.activation = activation()
+
+    def forward(self, x):
+        g1 = self.conv_1(x)
+        g2 = self.activation(self.conv_2(x))
+        return self.mixer(g1) * g2 + (1 - g2) * x
+
+
+class SmearGate(Layer):
+    """Smear gate (Modded-NanoGPT 2025-09 record): gate input-dependent + skip token-1.
+
+    Per ogni posizione t > 0:
+        x[t] += smear_lambda * sigmoid(W_gate @ x[t, :gate_input_dim]) * x[t-1]
+
+    Gate sparse (default 12 dim su d_model), input-dependent. Token 0 invariato.
+    smear_lambda init=0 ⇒ no-op iniziale. Receptive field = 1 (lookback=1 in fast gen).
+    """
+    __name__ = "SmearGate"
+
+    def __init__(self, d_model: int, gate_input_dim: int = 12):
+        super().__init__()
+        self.gate_input_dim = gate_input_dim
+        self.gate = nn.Linear(gate_input_dim, d_model, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+
+    def _init_weights(self):
+        # Identity-at-init: gate orthogonal; smear_lambda già 0 da __init__
+        # ⇒ shifted = 0 ⇒ no-op iniziale indipendentemente dal gate.
+        nn.init.orthogonal_(self.gate.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, L, D]. Caso L==1: nessuno "x[t-1]" disponibile → no-op.
+        if x.size(1) < 2:
+            return x
+        gate = torch.sigmoid(self.gate(x[:, 1:, :self.gate_input_dim]))  # [B, L-1, D]
+        prev = x[:, :-1, :]                                              # [B, L-1, D]
+        shifted = self.smear_lambda * gate * prev                        # [B, L-1, D]
+        return torch.cat([x[:, :1, :], x[:, 1:, :] + shifted], dim=1)
+
+
+class Embedder(Layer):
+    __name__ = "SymbolicEmbedder"
+    __complexity__ = "O(L d)"
+
+    def __init__(self, vocab_size, d_model: int):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.frozen = False
+        self.embedding = nn.Embedding(vocab_size, d_model)
+
+    def freeze(self):
+        if self.frozen:
+            flag("Trying to freeze an already frozen Embedder")
+        else:
+            self.embedding.weight.requires_grad = False
+            self.frozen = True
+
+    def unfreeze(self):
+        if self.frozen:
+            self.embedding.weight.requires_grad = True
+            self.frozen = False
+        else:
+            flag("Trying to unfreeze an already unfrozen Embedder")
+
+    def forward(self, x):
+        return self.embedding(x)
+
+    def finetune(self):
+        self.freeze()
+
+
+class EasyEmbedder(Layer):
+    __name__ = "SymbolicEmbedder"
+    __complexity__ = "O(L d)"
+
+    def __init__(self, vocab_size, d_model: int, dropout=0.13):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.dropout(self.embedding(x))
+
+
+class HybridAttentionBlock1d(Layer):
+    def __init__(self, d_model, sec_mixer, sec_params, attn_params, n_attn=1, n_sec=3,
+                 channel_mixer=MLP, channel_params=None):
+        super().__init__()
+        self.n_sec = n_sec
+        self.n_attn = n_attn
+        self.d_model = d_model
+        self.attn_params = attn_params
+        self.sec_params = sec_params
+        self.sec_mixer = sec_mixer
+        self.channel_mixer = channel_mixer
+        self.channel_params = channel_params
+
+        if channel_params is None and isinstance(channel_mixer, MLP):
+            self.channel_params = {'expand': 2}
+        else:
+            self.channel_params = channel_params
+        self.attn_blocks = nn.ModuleList([
+            Block1d(
+                channel_mixer=self.channel_mixer(d_model=d_model, **self.channel_params),
+                spatial_mixer=MultiheadAttentionMixer(d_model=d_model, **self.attn_params)
+            )
+            for _ in range(self.n_attn)
+        ])
+        self.sec_blocks = nn.ModuleList([
+            Block1d(
+                channel_mixer=self.channel_mixer(d_model=d_model, **self.channel_params),
+                spatial_mixer=self.sec_mixer(d_model=d_model, **self.sec_params)
+            )
+            for _ in range(self.n_sec)
+        ])
+
+    def forward(self, x):
+        for sec in self.sec_blocks:
+            x = sec(x)
+        for attn in self.attn_blocks:
+            x = attn(x)
+        return x
+
+
+########################################################################################################################
+#   VISION
+########################################################################################################################
+
+
+class PatchEmbedder(Layer):
+    __name__ = "PatchEmbedder"
+
+    def __init__(
+            self,
+            vocab_size=None,
+            d_model: int = 768,
+            img_size: Union[int, Tuple[int, int]] = 224,
+            patch_size: Union[int, Tuple[int, int]] = 16,
+            in_chans: int = 3,
+            flatten: bool = True,
+            norm_layer: Optional[Layer] = None,
+            cls: bool = False,
+    ):
+        super().__init__()
+        if isinstance(patch_size, int):
+            patch_size = (patch_size, patch_size)
+        if isinstance(img_size, int):
+            img_size = (img_size, img_size)
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+        self.embed_dim = d_model
+        self.flatten = flatten
+        self.frozen = False
+        self.norm = norm_layer if norm_layer is not None else None
+        self.use_cls = cls  # ← NEW
+
+        self.proj = nn.Conv2d(in_chans, d_model, kernel_size=patch_size, stride=patch_size)
+
+        if self.use_cls:
+            assert flatten, "CLS token requires flatten=True"
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))  # ← NEW
+
+        self.grid_size = None
+        if img_size is not None:
+            self.grid_size = (math.ceil(img_size[0] / patch_size[0]),
+                              math.ceil(img_size[1] / patch_size[1]))
+
+        self._num_patches = None
+        if self.grid_size is not None:
+            self._num_patches = self.grid_size[0] * self.grid_size[1]
+
+    def freeze(self):
+        if self.frozen:
+            flag("Trying to freeze an already frozen PatchEmbedder")
+            return
+        for p in self.proj.parameters():
+            p.requires_grad = False
+        if self.norm is not None:
+            for p in self.norm.parameters():
+                p.requires_grad = False
+        if self.use_cls:  # ← FREEZE CLS TOKEN TOO
+            self.cls_token.requires_grad = False
+        self.frozen = True
+
+    def unfreeze(self):
+        if not self.frozen:
+            flag("Trying to unfreeze an already unfrozen PatchEmbedder")
+            return
+        for p in self.proj.parameters():
+            p.requires_grad = True
+        if self.norm is not None:
+            for p in self.norm.parameters():
+                p.requires_grad = True
+        if self.use_cls:
+            self.cls_token.requires_grad = True
+        self.frozen = False
+
+    def num_patches(self, H: Optional[int] = None, W: Optional[int] = None) -> int:
+        if H is None or W is None:
+            if self.img_size is not None:
+                H, W = self.img_size
+            else:
+                raise ValueError("Must provide H and W or set img_size during init.")
+        ph, pw = self.patch_size
+        Hp = math.ceil(H / ph)
+        Wp = math.ceil(W / pw)
+        return Hp * Wp
+
+    def _pad_to_patch_multiple(self, x: torch.Tensor) -> torch.Tensor:
+        _, _, H, W = x.shape
+        ph, pw = self.patch_size
+        pad_h = (ph - H % ph) % ph
+        pad_w = (pw - W % pw) % pw
+        if pad_h == 0 and pad_w == 0:
+            return x
+        return nn.functional.pad(x, (0, pad_w, 0, pad_h), mode='constant', value=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        assert x.dim() == 4, "Input must be 4D (B, C, H, W)"
+        assert x.dtype in (torch.float32, torch.float16, torch.bfloat16), "Input dtype must be float"
+        B, C, H, W = x.shape
+        assert C == self.in_chans, f"Expected {self.in_chans} channels, got {C}"
+
+        x = self._pad_to_patch_multiple(x)
+
+        x = self.proj(x)  # (B, embed_dim, H_p, W_p)
+        H_p, W_p = x.shape[2], x.shape[3]
+        self._num_patches = H_p * W_p
+
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # (B, N, embed_dim)
+
+            if self.use_cls:
+                cls_tok = self.cls_token.expand(B, -1, -1)  # (B, 1, D)
+                x = torch.cat((cls_tok, x), dim=1)  # (B, 1+N, D)
+
+            if self.norm is not None:
+                x = self.norm(x)
+        else:
+            if self.norm is not None:
+                x = self.norm(x)
+
+        return x
+
+
+class MeanClassificationHead(Layer):
+    __name__ = "MeanClassificationHead"
+
+    def __init__(self, d_model, vocab_size):
+        super().__init__()
+        self.proj = nn.Linear(d_model, vocab_size)
+
+    def forward(self, x):
+        x = x.mean(dim=1)
+        return self.proj(x)
+
+
+class ClsHead(Layer):
+    __name__ = "Cls Head"
+
+    def __init__(self, d_model, vocab_size):
+        super().__init__()
+        self.proj = nn.Linear(d_model, vocab_size)
+
+    def forward(self, x):
+        x = x[:, -1:, :]
+        return self.proj(x)[:, 0, :]
+
+
+class MultiHeadUnembedder(Layer):
+    __name__ = "UnbiasedLinear"
+    __complexity__ = "O(L d^2)"
+
+    def __init__(self, d_model, vocab_size, k=4):
+        super(MultiHeadUnembedder, self).__init__()
+        self.linear = nn.Linear(d_model, vocab_size * k, bias=False)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class VathosModel(Layer):
+    __name__ = "VathosModel"
+
+    def __init__(self):
+        super().__init__()
+
+        self._losses = []
+        self._losses_dict = {}
+        self._losses_per_epoch = []
+        self._losses_per_epoch_dict = {}
+        self._losses_this_epoch = []
+        self._metrics = dict()
+        self._metrics_this_epoch = dict()
+        self._metrics_per_epoch = dict()
+        self.autosave = True
+        self.autosave_overwrite = True
+        self.best_loss = float('inf')
+        self.checkpoints = 0
+        self.steps = 0
+        self.steps_per_epoch = 0
+        self.epochs = 0
+        self.finetuning = False
+        self.name = 'VathosModel_defaultName'
+
+    def flag_not_training(self):
+        if not self.training:
+            flag("You are not training")
+
+    def register_loss(self, loss: float) -> bool:
+        """Registra loss + NaN/Inf detection.
+
+        Ritorna True se loss finita, False se NaN/Inf (con warning).
+        Mantiene `self._nan_streak` (consecutive NaN counter).
+        """
+        is_bad = (loss != loss) or (loss == float('inf')) or (loss == float('-inf'))
+        if is_bad:
+            self._nan_streak = getattr(self, "_nan_streak", 0) + 1
+            import warnings
+            warnings.warn(
+                f"VathosModel.register_loss: loss non-finita ({loss}) a step "
+                f"{self.steps}. nan_streak={self._nan_streak}.",
+                RuntimeWarning, stacklevel=2,
+            )
+        else:
+            self._nan_streak = 0
+        self._losses_dict[self.steps] = loss
+        self._losses.append(loss)
+        self._losses_this_epoch.append(loss)
+        self.steps += 1
+        return not is_bad
+
+    def get_last_loss(self):
+        return self._losses[-1]
+
+    @property
+    def nan_streak(self) -> int:
+        return getattr(self, "_nan_streak", 0)
+
+    def get_mean_loss(self, epoch=True):
+        if epoch:
+            return np.mean(self._losses_this_epoch)
+        else:
+            return np.mean(self._losses)
+
+    def register_metrics(self, metrics: dict):
+        for metric in metrics:
+            if metric in self._metrics:
+                self._metrics[metric].append(metrics[metric])
+                self._metrics_this_epoch[metric].append(metrics[metric])
+            else:
+                flag(f"Registering a new metric {metric}")
+                self._metrics[metric] = [metrics[metric]]
+                self._metrics_this_epoch[metric] = [metrics[metric]]
+
+    def register_epoch(self):
+        self.epochs += 1
+        self._losses_per_epoch.append(np.mean(self._losses_this_epoch))
+        self._losses_per_epoch_dict[self.steps] = np.mean(self._losses_this_epoch)
+        self._losses_this_epoch = []
+
+        for metric in self._metrics_this_epoch:
+            if metric in self._metrics_per_epoch:
+                self._metrics_per_epoch[metric].append(np.mean(self._metrics_this_epoch[metric]))
+            else:
+                self._metrics_per_epoch[metric] = [np.mean(self._metrics_this_epoch[metric])]
+            self._metrics_this_epoch[metric] = []
+
+        if self._losses_per_epoch[-1] < self.best_loss:
+            self.best_loss = self._losses_per_epoch[-1]
+            if self.autosave:
+                self.checkpoints += 1
+                if self.autosave_overwrite:
+                    self.save_checkpoint(f'{self.name}-checkpoint.pt')
+                else:
+                    self.save_checkpoint(f'{self.name}-checkpoint-{self.checkpoints}.pt')
+
+    def save_state_dict(self, path):
+        torch.save(self.state_dict(), path)
+
+    def plot_losses(self):
+        print(self.steps_per_epoch)
+        plt.plot(
+            list(self._losses_dict.keys()),
+            list(self._losses_dict.values()),
+            label="Losses", linewidth=1)
+        plt.plot(
+            list(self._losses_per_epoch_dict.keys()),
+            list(self._losses_per_epoch_dict.values()),
+            label="Losses Per Epoch", linewidth=2)
+
+        plt.xlabel("steps")
+        plt.ylabel("loss")
+        plt.title("Model Losses per Steps")
+        plt.show()
+
+    def plot_metrics(self, figsize=(12, 8)):
+        n_metrics = len(self._metrics_per_epoch)
+        n_plots = 1 + n_metrics
+
+        n_cols = 2
+        n_rows = (n_plots + 1) // 2
+
+        fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
+
+        if n_plots == 1:
+            axes = np.array([axes])
+        else:
+            axes = axes.flatten()
+
+        ax = axes[0]
+        ax.plot(
+            list(self._losses_dict.keys()),
+            list(self._losses_dict.values()),
+            label="Losses", linewidth=1, alpha=0.6)
+        ax.plot(
+            list(self._losses_per_epoch_dict.keys()),
+            list(self._losses_per_epoch_dict.values()),
+            label="Losses Per Epoch", linewidth=2)
+        ax.set_xlabel("Steps")
+        ax.set_ylabel("Loss")
+        ax.set_title("Training Loss")
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+
+        for idx, (metric_name, metric_values) in enumerate(self._metrics_per_epoch.items(), start=1):
+            ax = axes[idx]
+
+            if metric_name in self._metrics and len(self._metrics[metric_name]) > 0:
+                step_indices = list(range(len(self._metrics[metric_name])))
+                ax.plot(step_indices, self._metrics[metric_name],
+                        label=f"{metric_name}", linewidth=1, alpha=0.6)
+
+            epoch_indices = list(range(len(metric_values)))
+            ax.plot(epoch_indices, metric_values,
+                    label=f"{metric_name} Per Epoch", linewidth=2, marker='o')
+
+            ax.set_xlabel("Steps/Epochs")
+            ax.set_ylabel(metric_name)
+            ax.set_title(f"Metric: {metric_name}")
+            ax.legend()
+            ax.grid(True, alpha=0.3)
+
+        # Hide any unused subplots
+        for idx in range(n_plots, len(axes)):
+            axes[idx].set_visible(False)
+
+        plt.tight_layout()
+        plt.show()
+
+    def save_checkpoint(self, path):
+        """Save complete model checkpoint including training state"""
+        checkpoint = {
+            'model_state_dict': self.state_dict(),
+            'losses': self._losses,
+            'losses_dict': self._losses_dict,
+            'losses_per_epoch': self._losses_per_epoch,
+            'losses_per_epoch_dict': self._losses_per_epoch_dict,
+            'losses_this_epoch': self._losses_this_epoch,
+            'metrics': self._metrics,
+            'metrics_this_epoch': self._metrics_this_epoch,
+            'metrics_per_epoch': self._metrics_per_epoch,
+            'best_loss': self.best_loss,
+            'checkpoints': self.checkpoints,
+            'steps': self.steps,
+            'steps_per_epoch': self.steps_per_epoch,
+            'epochs': self.epochs,
+            'autosave': self.autosave,
+            'autosave_overwrite': self.autosave_overwrite,
+        }
+        torch.save(checkpoint, path)
+
+    def load_checkpoint(self, path):
+        """Load complete model checkpoint including training state"""
+        checkpoint = torch.load(path, weights_only=False)
+        self.load_state_dict(checkpoint['model_state_dict'])
+
+        # Restore training state
+        self._losses = checkpoint['losses']
+        self._losses_dict = checkpoint['losses_dict']
+        self._losses_per_epoch = checkpoint['losses_per_epoch']
+        self._losses_per_epoch_dict = checkpoint['losses_per_epoch_dict']
+        self._losses_this_epoch = checkpoint['losses_this_epoch']
+        self._metrics = checkpoint['metrics']
+        self._metrics_this_epoch = checkpoint['metrics_this_epoch']
+        self._metrics_per_epoch = checkpoint['metrics_per_epoch']
+        self.best_loss = checkpoint['best_loss']
+        self.checkpoints = checkpoint['checkpoints']
+        self.steps = checkpoint['steps']
+        self.steps_per_epoch = checkpoint['steps_per_epoch']
+        self.epochs = checkpoint['epochs']
+        self.autosave = checkpoint['autosave']
+        self.autosave_overwrite = checkpoint['autosave_overwrite']
+
+    def finetune(self):
+        self.finetuning = True
+
+    def train(self, *args, **kwargs):
+        self.finetuning = False
+        return super().train(*args, **kwargs)
+
+    def eval(self):
+        self.finetuning = False
+        return super().eval()
+
+
+########################################################################################################################
+#   Assemblers
+########################################################################################################################
+
+class SequenceModel(VathosModel):
+    __name__ = "SequenceModel"
+
+    def __init__(self, vocab_size: int, d_model: int, n_layers: int,
+                 max_len=1024,
+                 pos_encoder: bool | None | Layer | nn.Module = None,
+                 embedder=Embedder,
+                 embedder_args: dict = None,
+                 unembedder=UnbiasedLinear,
+                 unembedder_args=None,
+                 channel_mixer=MLP,
+                 spatial_mixer: Layer | nn.Module = MultiheadAttentionMixer,
+                 channel_args: dict = None,
+                 spatial_args: dict = None,
+                 name='',
+                 pad='none',
+                 baseblock=Block1d,
+                 baseblock_args=None,
+                 dropout=0.1,
+                 weight_tying=False,
+                 norm=nn.LayerNorm,
+                 d_modifiers: List | None = None,
+                 unet_skips=False,
+                 unet_weights=False
+                 ):
+        super().__init__()
+        self.pad = pad
+        flag(
+            "If you need to use any RoPE or alternative positional encodings which operate directly in the spatial mixer, be sure to call activate it in the spatial_args (e.g. rope=True)",
+            2)
+
+        if channel_args is None and channel_mixer is MLP:
+            channel_args = {"expand": 2, "activation": nn.GELU, "depth": 2}
+        if spatial_args is None:
+            spatial_args = {}
+        if channel_args is None:
+            channel_args = {}
+        if embedder_args is None:
+            embedder_args = {}
+        if unembedder_args is None:
+            unembedder_args = {'input_features': d_model, 'output_features': vocab_size}
+        if baseblock_args is None:
+            baseblock_args = {}
+        if d_modifiers is None:
+            d_modifiers = [1 for _ in range(d_model)]
+        else:
+            print(
+                f"{NUM} Initialized SequenceModel with d_model structure: {[int(d_model * d) for d in d_modifiers]}{RES}")
+
+        self.pipe = {}
+        self.name = name
+        self.baseblock = baseblock
+        self.spatial_mixer = spatial_mixer
+        self.channel_mixer = channel_mixer
+        self.baseblock_args = baseblock_args
+
+        self.spatial_args = spatial_args
+        self.channel_args = channel_args
+        self.vocab_size = vocab_size
+        self.max_len = max_len
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.unet_skips = unet_skips
+
+        if self.unet_skips:
+            self.skip_weights = nn.Parameter(torch.zeros(self.n_layers // 2), requires_grad=unet_weights)
+            if d_modifiers is not None:
+                flag("Unet Skips detected, assure d_modifier are symmetrical to make it work!")
+
+        self.embedder = embedder(vocab_size=vocab_size, d_model=d_model, **embedder_args)
+
+        self.pos_encoder = pos_encoder(d_model, max_len=max_len) if pos_encoder not in (True, False, None) else \
+            (SinusoidalPositionalEncoding(d_model, max_len=max_len) if pos_encoder is True else nn.Identity())
+
+        self.blocks = nn.ModuleList([
+            self.baseblock(
+                d_model=int(d_model * d_modifiers[i]),
+                channel_mixer=channel_mixer(d_model=int(d_model * d_modifiers[i]), **channel_args),
+                spatial_mixer=spatial_mixer(d_model=int(d_model * d_modifiers[i]), **spatial_args),
+                norm=norm,
+                **baseblock_args
+            )
+            for i in range(n_layers)
+        ])
+        self._piped_blocks = None
+
+        self.norm = nn.LayerNorm(d_model)
+        self.unembedder = unembedder(**unembedder_args)
+
+        if weight_tying and hasattr(self.embedder, 'embedding') and hasattr(self.unembedder, 'linear'):
+            self.unembedder.linear.weight = self.embedder.embedding.weight
+
+        elif weight_tying:
+            raise TypeError(
+                "Automatic weight tying is only possible if the the Embedder has a 'embedding' attribute and Unembedder has a linear attribute"
+                "You should manually do weight tying if you aim to use specific layer:"
+                "\n e.g. model.unembedder.yourmodule.weight = model.embedder.youembeddings.weight is an auto weight tying example")
+
+        self.embedder_complexity = embedder.__complexity__ if hasattr(embedder, "__complexity__") else "O(L d)"
+        self.unembedder_complexity = unembedder.__complexity__ if hasattr(unembedder, "__complexity__") else "O(L d)"
+        self.spatial_complextiy = spatial_mixer.__complexity__ if hasattr(spatial_mixer, "__complexity__") else "O(L d)"
+        self.channel_complexity = channel_mixer.__complexity__ if hasattr(channel_mixer, "__complexity__") else "O(L d)"
+        self.runned = False
+        self.embed_scale = math.sqrt(d_model)
+        self._init_weights()
+
+    def _init_weights(self):
+        """
+        Initialize weights following GPT-style conventions:
+        - Small std for embeddings
+        - Xavier/Kaiming for linear layers
+        - Scaled initialization for residual projections
+        """
+        std_embed = 0.02
+        try:
+            nn.init.normal_(self.embedder.embedding.weight, mean=0.0, std=std_embed)
+        except:
+            pass
+        try:
+            nn.init.normal_(self.unembedder.weight, mean=0.0, std=std_embed)
+        except:
+            pass
+
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                std_init = 0.02
+                nn.init.normal_(module.weight, mean=0.0, std=std_init)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+        for block_idx, block in enumerate(self.blocks):
+            for name, module in block.named_modules():
+                if isinstance(module, nn.Linear):
+                    depth_scale = (2.0 * self.n_layers) ** -0.5
+                    if 'out' in name.lower() or 'proj' in name.lower() or '.l2' in name or '.g2' in name:
+                        with torch.no_grad():
+                            module.weight.data *= depth_scale
+
+    def forward(self, x: torch.LongTensor, unembed=True):
+        # B, L = x.size(0), x.size(1)
+        # x = self.embedder(x) * self.embed_scale
+        x = self.pos_encoder(x)
+
+        skips = []
+        half_layers = self.n_layers // 2
+
+        for i, block in enumerate(self.blocks):
+            if self.unet_skips:
+                if i < half_layers:
+                    skips.append(x)
+                elif i >= self.n_layers - half_layers:
+                    skip_idx = self.n_layers - 1 - i
+                    x = x + self.skip_weights[skip_idx] * skips[skip_idx]
+
+            x = block(x)
+
+        if unembed:
+            x = self.norm(self.unembedder(x))
+
+        return x
+
+    def insert_block(self, idx, module):
+        self.blocks.insert(idx, module)
+
+    def append(self, module):
+        self.blocks.append(module)
+
+    @torch.no_grad()
+    def _clear_all_caches(self):
+        """Clear KV caches in all attention layers"""
+        for block in self.blocks:
+            for module in block.modules():
+                if hasattr(module, 'clear_cache'):
+                    module.clear_cache()
+
+    def forward(self, x: torch.LongTensor, unembed=True):
+        B, L = x.size(0), x.size(1)
+        x = self.embedder(x) * self.embed_scale
+        x = self.pos_encoder(x)
+
+        if self.pad == 'sqrt':
+            n = int((x.shape[1] ** 0.5) + 0.999999)
+            x = F.pad(x, (0, 0, 0, n ** 2 - L), mode="constant", value=0)
+
+        for block in self.blocks:
+            x = block(x)
+
+        if unembed:
+            x = self.unembedder(x)
+
+        return x[:, :L, :]
+
+    def _sample_token(self, logits, temperature=1.0, top_p=1.0, top_k=None):
+        logits = logits / (temperature + 1e-8)
+
+        if top_k is not None and top_k > 0:
+            top_k = min(top_k, logits.size(-1))
+            v, _ = torch.topk(logits, top_k)
+            logits[logits < v[:, [-1]]] = float('-inf')
+
+        if top_p < 1.0:
+            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+            sorted_indices_to_remove = cumulative_probs > top_p
+            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+            sorted_indices_to_remove[..., 0] = 0
+
+            indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+            logits[indices_to_remove] = float('-inf')
+
+        probs = F.softmax(logits, dim=-1)
+        return torch.multinomial(probs, num_samples=1)
+
+    def generate(self, *args, **kwargs):
+        if not 'custom_generate' in kwargs:
+            kwargs['custom_generate'] = False
+        if kwargs['custom_generate']:
+            del kwargs['custom_generate']
+            return self.custom_generate(*args, **kwargs)
+        else:
+            del kwargs['custom_generate']
+            return self.simple_generate(*args, **kwargs)
+
+    @torch.no_grad()
+    def simple_generate(self, prompt: torch.Tensor, max_len=100, temperature=1.0,
+                        top_p=1.0, top_k=50, token_end=None, repetition_penalty=1.0):
+        self.eval()
+        if prompt.dim() == 1:
+            prompt = prompt.unsqueeze(0)
+
+        generated = prompt.clone()
+        pbar = tqdm(range(max_len), desc="Simple Gen")
+
+        for _ in pbar:
+            logits = self.forward(generated, unembed=True)
+            next_token_logits = logits[:, -1, :]
+
+            # Apply repetition penalty
+            if repetition_penalty != 1.0:
+                next_token_logits = self._apply_repetition_penalty(
+                    next_token_logits, generated, repetition_penalty
+                )
+
+            next_token = self._sample_token(next_token_logits, temperature, top_p, top_k)
+            generated = torch.cat([generated, next_token], dim=1)
+
+            if token_end is not None and (next_token == token_end).all():
+                break
+
+        return generated
+
+    @torch.no_grad()
+    def custom_generate(self, prompt: torch.Tensor, max_len=100, temperature=1.0,
+                        top_p=1.0, top_k=50, token_end=None, repetition_penalty=1.0):
+        self.eval()
+        self._clear_all_caches()
+
+        if prompt.dim() == 1:
+            prompt = prompt.unsqueeze(0)
+
+        x = self.embedder(prompt) * self.embed_scale
+        if self.pos_encoder is not None and not isinstance(self.pos_encoder, nn.Identity):
+            x = self.pos_encoder(x)
+
+        for block in self.blocks:
+            if block.has_custom_generate():
+                x = block.generate(x)
+            else:
+                x = block(x)
+
+        generated = prompt.clone()
+        logits = self.unembedder(x)
+        next_token_logits = logits[:, -1, :]
+
+        # Apply repetition penalty
+        if repetition_penalty != 1.0:
+            next_token_logits = self._apply_repetition_penalty(
+                next_token_logits, generated, repetition_penalty
+            )
+
+        next_token = self._sample_token(next_token_logits, temperature, top_p, top_k)
+        generated = torch.cat([generated, next_token], dim=1)
+
+        pbar = tqdm(range(max_len - 1), desc="Fast Gen")
+        for _ in pbar:
+            x_t = self.embedder(next_token) * self.embed_scale
+            current_pos = generated.shape[1] - 1
+
+            if isinstance(self.pos_encoder, SinusoidalPositionalEncoding):
+                pe_slice = self.pos_encoder.pe[current_pos: current_pos + 1].unsqueeze(0)
+                x_t = x_t + pe_slice
+
+            for block in self.blocks:
+                if block.has_custom_generate():
+                    x_t = block.generate(x_t)
+                else:
+                    x_t = block(x_t)
+
+            logits = self.unembedder(x_t)
+            next_token_logits = logits[:, -1, :]
+
+            # Apply repetition penalty
+            if repetition_penalty != 1.0:
+                next_token_logits = self._apply_repetition_penalty(
+                    next_token_logits, generated, repetition_penalty
+                )
+
+            next_token = self._sample_token(next_token_logits, temperature, top_p, top_k)
+            generated = torch.cat([generated, next_token], dim=1)
+
+            if token_end is not None and (next_token == token_end).all():
+                break
+
+        self._clear_all_caches()
+        return generated
+
+    def _apply_repetition_penalty(self, logits: torch.Tensor,
+                                  generated: torch.Tensor,
+                                  repetition_penalty: float) -> torch.Tensor:
+        """
+        Apply repetition penalty to logits based on previously generated tokens.
+
+        Args:
+            logits: Shape (batch_size, vocab_size)
+            generated: Shape (batch_size, seq_len) - previously generated tokens
+            repetition_penalty: Penalty factor (> 1.0 discourages repetition)
+
+        Returns:
+            Modified logits with repetition penalty applied
+        """
+        batch_size = logits.shape[0]
+
+        for i in range(batch_size):
+            # Get unique tokens in the generated sequence for this batch item
+            unique_tokens = generated[i].unique()
+
+            # Apply penalty: divide logits by penalty if positive, multiply if negative
+            for token_id in unique_tokens:
+                if logits[i, token_id] > 0:
+                    logits[i, token_id] /= repetition_penalty
+                else:
+                    logits[i, token_id] *= repetition_penalty
+
+        return logits
+
+    def summary(self):
+        complexities = [self.channel_complexity, self.spatial_complextiy, self.embedder_complexity,
+                        self.unembedder_complexity]
+        print(f'{NUM}VATHOS{RES} {self.name} Summary:')
+        print(f"{NUM}SequenceModel{RES}(d_model={NUM}{self.d_model}{RES}, n_layer={NUM}{self.n_layers}{RES})")
+        print(f"\t - {NUM}VOCAB_SIZE:{RES}: {NUM}{self.vocab_size}{RES}")
+        print(f"\t - {NUM}D_MODEL:{RES}: {NUM}{self.d_model}{RES}")
+        print(f"\t - {NUM}N_LAYERS:{RES}: {NUM}{self.n_layers}{RES}")
+        print("")
+        print(f"\t - {NUM}Embedder{RES}: {getname(self.embedder)} - {NUM}{self.embedder_complexity}{RES}")
+        print(f"\t - {NUM}Unembedder{RES}: {getname(self.unembedder)} - {NUM}{self.unembedder_complexity}{RES}")
+        print(
+            f"\t - {NUM}Spatial Mixer{RES}: {getname(self.spatial_mixer)}({self.spatial_args}) - {NUM}{self.spatial_complextiy}{RES}")
+        print(
+            f"\t - {NUM}Channel Mixer{RES}: {getname(self.channel_mixer)}({self.channel_args}) - {NUM}{self.channel_complexity}{RES}")
+        print(f"Num Parameters: {NUM}{sum([p.numel() for p in self.parameters()]):_}{RES}")
+        print(f"Num Trainable Parameters: {NUM}{sum([p.numel() for p in self.parameters() if p.requires_grad]):_}{RES}")
+        print(f"Total Complexity: {NUM}{combine_big_o_sum(complexities)}{RES}")
+
+    def finetune(self):
+        flag(
+            "Finetune simply checks for finetune() methods in spatial mixers, channel mixers, embedder and unembedder, "
+            "if a finetune method is not available, the module/Layer will be left as it is")
+        super().finetune()
+
+        if hasattr(self.embedder, "finetune"):
+            self.embedder.finetune()
+        if hasattr(self.unembedder, "finetune"):
+            self.unembedder.finetune()
+        for block in self.blocks:
+            if hasattr(block, "finetune"):
+                block.channel_mixer.finetune()
+                block.spatial_mixer.finetune()
+
+
+class ModdedFormer(VathosModel):
+    def __init__(self, vocab_size: int, embed_dim: int, d_models: List[int], spatials: List[Builder], expand=3,
+                 M_dims: List[int] | None = None, weights_tying=True,
+                 baseblock=Block1d, norm=RMSNorm, ffn_act=ReLU2, unet_skips=False, max_len=2400, zeroskip=False,
+                 UDLP=VariableUDLP,
+                 skips: List[int | None] | None = None, value_embeddings: List[int | None] | None = None,
+                 ve_type: str = 'scalar', ve_gate_dim: int | None = None,
+                 learnable_pe: bool = False, input_projections: bool | int=False, smear_gate: None | nn.Module | Layer=False):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        assert d_models[0] == embed_dim
+        assert max(d_models) == min(d_models), "Variables d_models is WIP and not working"
+        self.d_models = d_models
+        n_layers = self.n_layer = len(d_models)
+        self.input_projections = input_projections
+        self.smear_gate = smear_gate
+        if M_dims is not None:
+            self.M_dims = M_dims
+        else:
+            self.M_dims = [d * expand for d in d_models]
+
+        self.embedder = Embedder(vocab_size, embed_dim)
+        self.unembedder = UnbiasedLinear(d_models[-1], vocab_size)
+        self.spatials = spatials
+        self.norm = norm
+        self.unet = unet_skips
+        self.max_len = max_len
+        self.zeroskip = zeroskip
+
+        self.skips = skips if skips is not None else [None] * n_layers
+        assert len(self.skips) == n_layers
+
+        self.skip_lambdas = nn.ParameterDict()
+        for source_idx, target_idx in enumerate(self.skips):
+            if target_idx is not None:
+                assert target_idx > source_idx
+                assert target_idx < n_layers
+                self.skip_lambdas[f"route_{source_idx}_to_{target_idx}"] = nn.Parameter(torch.zeros(1))
+
+        if weights_tying:
+            assert d_models[-1] == embed_dim
+            self.unembedder.linear.weight = self.embedder.embedding.weight
+
+        blocks = []
+        for i in range(n_layers):
+            in_dim = embed_dim if i == 0 else self.d_models[i - 1]
+            out_dim = self.d_models[i]
+            blocks.append(
+                baseblock(
+                    self.d_models[i],
+                    UDLP(in_dim, d_output=out_dim, M=self.M_dims[i], activation=ffn_act),
+                    self.spatials[i](self.d_models[i]),
+                    norm=norm
+                ))
+
+        self.blocks = nn.ModuleList(blocks)
+
+        if zeroskip:
+            self.zeroskip_params = nn.ParameterList([nn.Parameter(torch.tensor([0.0])) for _ in range(n_layers)])
+        if input_projections > 0:
+            if not zeroskip:
+                raise ValueError("Zeroskip must be activated when using input_projections")
+            self.inputs_projection_linears = nn.Linear(in_features=d_models[0], out_features=d_models[0]*input_projections, bias=False)
+            self.inputs_projection_params = nn.ParameterList([nn.Parameter(torch.tensor([0.0])) for _ in range(n_layers) for i in range(input_projections)])
+
+
+        assert ve_type in ('scalar', 'gate'), f"ve_type must be 'scalar' or 'gate', got {ve_type}"
+        self.ve_type = ve_type
+        self.ve_gate_dim = ve_gate_dim if ve_gate_dim is not None else embed_dim
+        self.value_embeddings_cfg = value_embeddings
+
+        # Value embeddings: ModuleDict/ParameterDict come prima (per leggibilità + checkpoint legacy),
+        # MA in più costruiamo strutture INDEXED-BY-LAYER per il forward (no str() runtime, no dict
+        # lookup → compile-friendly).
+        if value_embeddings is not None:
+            assert len(value_embeddings) == n_layers, "value_embeddings length must equal n_layers"
+            unique_groups = sorted(set(v for v in value_embeddings if v is not None))
+            self.ve_embeddings = nn.ModuleDict({
+                str(g): nn.Embedding(vocab_size, embed_dim)
+                for g in unique_groups
+            })
+            if ve_type == 'scalar':
+                self.ve_scales = nn.ParameterDict({
+                    str(i): nn.Parameter(torch.zeros(1))
+                    for i, v in enumerate(value_embeddings) if v is not None
+                })
+            else:  # gate
+                self.ve_gates = nn.ModuleDict({
+                    str(i): nn.Linear(self.ve_gate_dim, embed_dim, bias=False)
+                    for i, v in enumerate(value_embeddings) if v is not None
+                })
+
+            # Strutture parallele indexed-by-layer per il forward (riferiscono gli stessi param).
+            self._ve_per_layer = nn.ModuleList([
+                self.ve_embeddings[str(v)] if v is not None else nn.Identity()
+                for v in value_embeddings
+            ])
+            self._has_ve_per_layer = tuple(v is not None for v in value_embeddings)
+            self._has_any_ve = any(self._has_ve_per_layer)
+
+            if ve_type == 'scalar':
+                self._ve_scale_per_layer = nn.ParameterList([
+                    self.ve_scales[str(i)] if value_embeddings[i] is not None
+                    else nn.Parameter(torch.zeros(1), requires_grad=False)
+                    for i in range(n_layers)
+                ])
+                self._ve_gate_per_layer = nn.ModuleList([nn.Identity() for _ in range(n_layers)])
+            else:
+                self._ve_scale_per_layer = nn.ParameterList([
+                    nn.Parameter(torch.zeros(1), requires_grad=False) for _ in range(n_layers)
+                ])
+                self._ve_gate_per_layer = nn.ModuleList([
+                    self.ve_gates[str(i)] if value_embeddings[i] is not None else nn.Identity()
+                    for i in range(n_layers)
+                ])
+        else:
+            self.ve_embeddings = nn.ModuleDict()
+            self._has_ve_per_layer = tuple([False] * n_layers)
+            self._has_any_ve = False
+            self._ve_per_layer = nn.ModuleList([nn.Identity() for _ in range(n_layers)])
+            self._ve_scale_per_layer = nn.ParameterList()
+            self._ve_gate_per_layer = nn.ModuleList()
+
+        # ---------------------------------------------------------------------
+        # Skip routing: pre-computed plan + ParameterList (no string keys, no dict mutato)
+        # ---------------------------------------------------------------------
+        # Per ogni layer i:
+        #   - _skip_save[i]: bool, se salvare l'output come sorgente di skip
+        #   - _skip_consume_at[i]: tuple di (source_idx, gate_idx_in_list) da consumare qui
+        self._skip_save = tuple(s is not None for s in self.skips)
+        consume_lists = [[] for _ in range(n_layers)]
+        gate_param_count = 0
+        # Mapping (source, target) -> gate index nella ParameterList
+        self._skip_gate_map = {}
+        for source_idx, target_idx in enumerate(self.skips):
+            if target_idx is not None:
+                self._skip_gate_map[(source_idx, target_idx)] = gate_param_count
+                consume_lists[target_idx].append((source_idx, gate_param_count))
+                gate_param_count += 1
+        self._skip_consume_at = tuple(tuple(lst) for lst in consume_lists)
+        self._has_any_skip = gate_param_count > 0
+        # ParameterList parallela: stessi tensori della ParameterDict legacy, indicizzati posizionalmente.
+        if self._has_any_skip:
+            self._skip_gates_list = nn.ParameterList([
+                self.skip_lambdas[f"route_{s}_to_{t}"]
+                for (s, t), _ in sorted(self._skip_gate_map.items(), key=lambda kv: kv[1])
+            ])
+        else:
+            self._skip_gates_list = nn.ParameterList()
+
+        self.learnable_pe = learnable_pe
+        if self.learnable_pe:
+            self.pos_emb = nn.Parameter(torch.zeros(1, max_len, embed_dim))
+
+        self.final_norm = RMSNorm(d_models[-1])
+        self._init_weights()
+
+        # Flag globale fast-path: nessuna feature opzionale → loop pulito senza branch.
+        self._has_any_zeroskip = bool(zeroskip)
+        self._has_input_projections = int(input_projections) > 0
+        self._has_smear_gate = bool(smear_gate)
+        self._has_learnable_pe = bool(learnable_pe)
+        self._fast_path = not any([
+            self._has_any_ve, self._has_any_skip,
+            self._has_any_zeroskip, self._has_input_projections,
+            self._has_smear_gate, self._has_learnable_pe,
+        ])
+
+    def _apply_ve_weight(self, ve: torch.Tensor, x: torch.Tensor, layer_idx: int) -> torch.Tensor:
+        """Legacy helper (mantenuto per backward compat). Il forward usa la versione indicizzata."""
+        if self.ve_type == 'scalar':
+            return ve * self.ve_scales[str(layer_idx)]
+        else:
+            gate = 2.0 * torch.sigmoid(self.ve_gates[str(layer_idx)](x[..., :self.ve_gate_dim]))
+            return ve * gate
+
+    def _compute_ve_indexed(self, input_ids: torch.Tensor, x: torch.Tensor, i: int):
+        """Versione fast: ModuleList[i] invece di ModuleDict[str(group_id)]."""
+        ve = self._ve_per_layer[i](input_ids)
+        if self.ve_type == 'scalar':
+            return ve * self._ve_scale_per_layer[i]
+        # gate type
+        return ve * (2.0 * torch.sigmoid(self._ve_gate_per_layer[i](x[..., :self.ve_gate_dim])))
+
+    # -------------------------------------------------------------------------
+    # Forward — fast path se niente feature opzionali, slow path altrimenti
+    # -------------------------------------------------------------------------
+
+    def forward(self, x):
+        # FAST PATH: standard transformer pretrain, no skip / no VE / no zeroskip / no PE / no smear.
+        # È il caso di gran lunga più comune (es. PiCO 2 / Guido-1 baseline).
+        if self._fast_path:
+            x = self.embedder(x)
+            for block in self.blocks:
+                x = block(x)
+            return self.unembedder(self.final_norm(x))
+
+        # SLOW PATH: include tutte le feature opzionali. Anche qui niente string keys / dict mutati:
+        # le strutture sono pre-computate al __init__ e indexed-by-layer.
+        input_ids = x  # alias per VE lookup
+
+        x0 = self.embedder(x)
+        if self._has_input_projections:
+            xs = self.inputs_projection_linears(x0).chunk(self.input_projections, dim=-1)
+        if self._has_learnable_pe:
+            x0 = x0 + self.pos_emb[:, :x.size(1), :]
+        if self._has_smear_gate:
+            x0 = self.smear_gate(x0)
+
+        x = x0
+        # Buffer di skip: lista Python di lunghezza fissa = n_layers. Allocata solo se serve.
+        skip_buffer = [None] * self.n_layer if self._has_any_skip else None
+
+        for i, block in enumerate(self.blocks):
+            # VE — indexed lookup (no string)
+            ve = None
+            if self._has_ve_per_layer[i]:
+                ve = self._compute_ve_indexed(input_ids, x, i)
+
+            # Block forward + zeroskip variants
+            if self._has_any_zeroskip and self._has_input_projections:
+                proj_params = self.inputs_projection_params[
+                    i * self.input_projections: (i + 1) * self.input_projections
+                ]
+                x = block(x, ve=ve) + x0 * self.zeroskip_params[i] + sum(
+                    [xj * p for xj, p in zip(xs, proj_params)]
+                )
+            elif self._has_any_zeroskip:
+                x = block(x, ve=ve) + x0 * self.zeroskip_params[i]
+            else:
+                x = block(x, ve=ve)
+
+            # Skip routing — pre-computed plan, indexed gates (no dict, no string)
+            if self._has_any_skip:
+                consume = self._skip_consume_at[i]
+                for source_idx, gate_idx in consume:
+                    x = x + self._skip_gates_list[gate_idx] * skip_buffer[source_idx]
+                if self._skip_save[i]:
+                    skip_buffer[i] = x
+
+        return self.unembedder(self.final_norm(x))
+
+    @torch.no_grad()
+    def _clear_all_caches(self):
+        for block in self.blocks:
+            for module in block.modules():
+                if hasattr(module, 'clear_cache'):
+                    module.clear_cache()
+
+    def _embed_input(self, ids: torch.Tensor, pos_start: int) -> torch.Tensor:
+        x0 = self.embedder(ids)
+        if self.learnable_pe:
+            L = ids.size(1)
+            x0 = x0 + self.pos_emb[:, pos_start:pos_start + L, :]
+        if self.smear_gate:
+            x0 = self.smear_gate(x0)
+        return x0
+
+    def _blocks_generate(self, x: torch.Tensor, x0: torch.Tensor,
+                        xs, input_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Mirror of forward()'s block loop, but calls block.generate(...) so
+        that attention mixers populate / consume their KV cache.
+        Works for both prefill (L>1) and single-token increments (L=1).
+        """
+        active_skips = {}
+        for i, block in enumerate(self.blocks):
+            ve = None
+            if self.value_embeddings_cfg is not None and self.value_embeddings_cfg[i] is not None:
+                group_id = self.value_embeddings_cfg[i]
+                ve = self.ve_embeddings[str(group_id)](input_ids)
+                ve = self._apply_ve_weight(ve, x, i)
+
+            if self.zeroskip and self.input_projections > 0:
+                proj_params = self.inputs_projection_params[
+                    i * self.input_projections: (i + 1) * self.input_projections
+                ]
+                x = block.generate(x, ve=ve) + x0 * self.zeroskip_params[i] + sum(
+                    [xj * p for xj, p in zip(xs, proj_params)]
+                )
+            elif self.zeroskip:
+                x = block.generate(x, ve=ve) + x0 * self.zeroskip_params[i]
+            else:
+                x = block.generate(x, ve=ve)
+
+            for source_idx, target_idx in enumerate(self.skips):
+                if target_idx == i:
+                    gate = self.skip_lambdas[f"route_{source_idx}_to_{target_idx}"]
+                    x = x + gate * active_skips[source_idx]
+                    del active_skips[source_idx]
+
+            if self.skips[i] is not None:
+                active_skips[i] = x
+        return x
+
+    @torch.no_grad()
+    def generate(self, prompt: torch.Tensor, max_len: int = 100,
+                 temperature: float = 1.0, top_k: int | None = None,
+                 top_p: float = 1.0, repetition_penalty: float = 1.0,
+                 token_end=None, simple: bool = False):
+        """
+        KV-cached AR generation. Prefills the prompt in one shot (every spatial
+        mixer's .generate accepts L>1 and populates its kv_cache), then walks
+        token-by-token reusing the cached keys/values.
+
+        simple=True falls back to simple_ar_generate (O(L^2)) — useful for
+        regression-checking the cached path.
+        """
+        if simple:
+            return simple_ar_generate(
+                self, prompt, max_len=max_len, temperature=temperature,
+                top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty,
+                token_end=token_end,
+            )
+
+        self.eval()
+        self._clear_all_caches()
+
+        if self.value_embeddings_cfg is not None:
+            flag("ModdedFormer.generate: value_embeddings only work with spatial mixers "
+                 "whose .generate accepts a 've=' kwarg (e.g. MultiheadAttentionMixer). "
+                 "On other mixers ve is silently dropped during generation.", 2)
+
+        if prompt.dim() == 1:
+            prompt = prompt.unsqueeze(0)
+        device = next(self.parameters()).device
+        prompt = prompt.to(device)
+
+        L = prompt.size(1)
+        x0 = self._embed_input(prompt, pos_start=0)
+        xs = (self.inputs_projection_linears(x0).chunk(self.input_projections, dim=-1)
+              if self.input_projections > 0 else None)
+
+        x = self._blocks_generate(x0, x0, xs, prompt)
+        logits = self.unembedder(self.final_norm(x))[:, -1, :]
+
+        if repetition_penalty != 1.0:
+            logits = apply_repetition_penalty(logits, prompt, repetition_penalty)
+        next_token = sample_next_token(logits, temperature, top_k, top_p)
+        generated = torch.cat([prompt, next_token], dim=1)
+
+        pos = L
+        pbar = tqdm(range(max_len - 1), desc="ModdedFormer fast gen")
+        for _ in pbar:
+            x0_t = self._embed_input(next_token, pos_start=pos)
+            xs_t = (self.inputs_projection_linears(x0_t).chunk(self.input_projections, dim=-1)
+                    if self.input_projections > 0 else None)
+
+            x_t = self._blocks_generate(x0_t, x0_t, xs_t, next_token)
+            logits = self.unembedder(self.final_norm(x_t))[:, -1, :]
+
+            if repetition_penalty != 1.0:
+                logits = apply_repetition_penalty(logits, generated, repetition_penalty)
+            next_token = sample_next_token(logits, temperature, top_k, top_p)
+            generated = torch.cat([generated, next_token], dim=1)
+            pos += 1
+
+            if token_end is not None and (next_token == token_end).all():
+                break
+
+        self._clear_all_caches()
+        return generated
+
+    def _init_weights(self):
+        std_embed = 0.02
+        try:
+            nn.init.normal_(self.embedder.embedding.weight, mean=0.0, std=std_embed)
+        except:
+            pass
+        try:
+            nn.init.normal_(self.unembedder.weight, mean=0.0, std=std_embed)
+        except:
+            pass
+
+        if self.learnable_pe:
+            nn.init.normal_(self.pos_emb, mean=0.0, std=std_embed)
+
+        for block_idx, block in enumerate(self.blocks):
+            block.channel_mixer._init_weights()
+            try:
+                block.spatial_mixer._init_weights()
+            except:
+                pass
+
+    def summary(self):
+        print(f"Vathos {NUM}ModdedFormer{RES} Summary")
+        print(f"Embedding dim: {self.embed_dim}")
+        print(f"Learnable PE: {self.learnable_pe}")
+        print(f"Skips: {self.skips}")
+        for i in range(self.n_layer):
+            if i == 0:
+                in_dim = self.embed_dim
+            else:
+                in_dim = self.d_models[i - 1]
+            out_dim = self.d_models[i]
+            M = self.M_dims[i]
+            print(f"Layer {i}: {in_dim} -> {out_dim}")
+            print(f"Block: {i}: {self.blocks[i]}")
+            print(f"\tFFN Dimension: {M}")
+            print(f"\tFFN Activation: {self.blocks[i].channel_mixer.activation}")
+
+        print(f"Num Parameters: {NUM}{sum([p.numel() for p in self.parameters()]):_}{RES}")
+        print(f"Num Trainable Parameters: {NUM}{sum([p.numel() for p in self.parameters() if p.requires_grad]):_}{RES}")
+
+
+#######################################################################################################################
+
+"""def test_causality(module=MTransformer(8, 4, 2)):
+    torch.manual_seed(42)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = module
+    x = torch.randn(2, 20, model.d_model, device=device)
+
+    print(f"\nInput shape: {x.shape}")
+
+    with torch.no_grad():
+        output = model(x)
+
+    model.eval()
+    with torch.no_grad():
+        full_output = model(x)
+
+        for t in range(1, 20):
+            prefix_output = model(x[:, :t, :])
+
+            max_diff = (full_output[:, :t, :] - prefix_output).abs().max().item()
+
+            if max_diff > 1e-5:
+                print(f"Causality violated")
+                break
+        else:
+            print("Causality verified")
+
+
+def test_causality_symbolic(module=SequenceModel(128, 16, 4, 2, pos_encoder=True)):
+    torch.manual_seed(42)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = module
+    x = torch.randint(0, model.vocab_size, (2, model.max_len), device=device)
+
+    print(f"\nInput shape: {x.shape}")
+
+    with torch.no_grad():
+        output = model(x)
+
+    model.eval()
+    with torch.no_grad():
+        full_output = model(x, unembed=False)
+
+        for t in range(1, 20):
+            prefix_output = model(x[:, :t], unembed=False)
+
+            max_diff = (full_output[:, :t] - prefix_output).abs().max().item()
+
+            if max_diff > 1e-5:
+                print(f"Causality violated")
+                break
+        else:
+            print("Causality verified")
+
+
+def test_symbolic_model(model):
+    vocab_size = model.vocab_size
+    length = model.max_len
+    x = torch.randint(0, vocab_size, (2, length))
+    print(f"Input shape {x.shape}")
+    x = model(x)
+    print(f"Output shape {x.shape}")
+    print(f"Bounds:  min:{x.min()}, max:{x.max()}, mean:{x.mean()}, std:{x.std()}, sum_of_a_vector: {x[0, 0, :].sum()}")"""
+
+########################################################################################################################
+
+NAMES = {
+    "MLP": MLP,
+    "Attention": MultiheadAttentionMixer,
+    "MHA": MultiheadAttentionMixer,
+    "Embed": Embedder,
+    "EMBED": Embedder,
+    "E": Embedder,
+    "PE": SinusoidalPositionalEncoding,
+    "PatchEmbedder": PatchEmbedder,
+    "ClassificationHead": MeanClassificationHead,
+    "MCH": MeanClassificationHead,
+    "ClsHead": ClsHead,
+}
+
+
+def expand_architecture(arch_string):
+    pattern = r"(\(.+?\))x(\d+)"
+
+    def replacer(match):
+        content = match.group(1)
+        count = int(match.group(2))
+        return " -> ".join([content] * count)
+
+    return re.sub(pattern, replacer, arch_string)
+
+
+def assemble(code, d_model=512):
+    code = code.strip()
+    if "x" in code:
+        code = expand_architecture(code)
+    divided = code.split("->")
+
+    for arch in divided:
+        arch = arch.strip()
+        if arch.startswith("("):
+            layers = []
+            for subarch in arch[1:-1].split(","):
+                subarch = subarch.strip()
+                if subarch in NAMES:
+                    layer = NAMES[subarch](d_model=d_model)
+                else:
+                    raise ValueError(f"Unknown Layer: {subarch}")
+                layers.append(subarch)
+            block = Block1d(*layers)
+
+        else:
+            if arch in NAMES:
+                layer = NAMES[arch]
+            elif hasattr(nn, arch):
+                layer = getattr(nn, arch)
+            else:
+                raise ValueError(f"Unknown Layer: {arch}")
+
+
+def wrap(module):
+    return tWrapper(module)
+
+
+def get_builder(layer, params):
+    pass
+
+
+########################################################################################################################
+# Pre Built
+########################################################################################################################
+
+
+GQA = GroupedQueryAttention
+GQANOO = GroupedQueryAttentionNOO
+GQANOV = GroupedQueryAttentionNOV
+Attention = MultiheadAttentionMixer
+AttentionNOV = MultiheadAttentionMixerNOV
+NOVLa2 = MultiheadAttentionMixerNOVLa2
+
+
+# =============================================================================
+# HIGH-LEVEL MODELS
+# =============================================================================
+#
+# Modelli pronti all'uso, ottimizzati per training/inference efficienti. A
+# differenza di SequenceModel / ModdedFormer (pensati per studio + ablation),
+# qui ogni Python branch è pre-risolto al __init__, niente ParameterDict /
+# ModuleDict string-keyed, niente dict mutati nel forward. Tutto torch.compile-
+# friendly.
+#
+# d_model è fissa su tutti i layer (sceltà per semplicità + compile efficiency).
+# =============================================================================
+
+
+class BlockX0(Layer):
+    """Pre-norm block + x0-skip nativo.
+
+    Forward(x, x0) -> x + spatial(norm1(x)) + channel(norm2(x)) + x0_lambda * x0
+    `x0_lambda` init=0 ⇒ comportamento identico a Block1d al primo step.
+    """
+    __name__ = "BlockX0"
+
+    def __init__(self, d_model: int, channel_mixer: Layer, spatial_mixer: Layer, norm=RMSNorm):
+        super().__init__()
+        self.spatial_mixer = spatial_mixer
+        self.channel_mixer = channel_mixer
+        self.norm1 = norm(d_model)
+        self.norm2 = norm(d_model)
+        self.x0_lambda = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x: torch.Tensor, x0: torch.Tensor) -> torch.Tensor:
+        x = x + self.spatial_mixer(self.norm1(x))
+        x = x + self.channel_mixer(self.norm2(x))
+        return x + self.x0_lambda * x0
+
+    def generate(self, x: torch.Tensor, x0: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        sm = self.spatial_mixer
+        x = x + (sm.generate(h) if sm.has_custom_generate() else sm(h))
+        h = self.norm2(x)
+        cm = self.channel_mixer
+        x = x + (cm.generate(h) if cm.has_custom_generate() else cm(h))
+        return x + self.x0_lambda * x0
+
+
+class PiCOBlock(BlockX0):
+    """BlockX0 + value-embedding nativo (off di default, zero costo se off).
+
+    Quando `ve_enabled=True`, forward accetta un tensor `ve` esterno e lo passa
+    al spatial_mixer pre-scalato per `ve_scale` (Parameter init=0 ⇒ no-op iniziale).
+    """
+    __name__ = "PiCOBlock"
+
+    def __init__(self, d_model: int, channel_mixer: Layer, spatial_mixer: Layer,
+                 norm=RMSNorm, ve_enabled: bool = False):
+        super().__init__(d_model, channel_mixer, spatial_mixer, norm)
+        self.ve_enabled = ve_enabled
+        if ve_enabled:
+            self.ve_scale = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x: torch.Tensor, x0: torch.Tensor, ve: torch.Tensor = None) -> torch.Tensor:
+        if self.ve_enabled and ve is not None:
+            x = x + self.spatial_mixer(self.norm1(x), ve=ve * self.ve_scale)
+        else:
+            x = x + self.spatial_mixer(self.norm1(x))
+        x = x + self.channel_mixer(self.norm2(x))
+        return x + self.x0_lambda * x0
+
+    def generate(self, x: torch.Tensor, x0: torch.Tensor, ve: torch.Tensor = None) -> torch.Tensor:
+        h = self.norm1(x)
+        sm = self.spatial_mixer
+        if self.ve_enabled and ve is not None:
+            scaled = ve * self.ve_scale
+            x = x + (sm.generate(h, ve=scaled) if sm.has_custom_generate() else sm(h, ve=scaled))
+        else:
+            x = x + (sm.generate(h) if sm.has_custom_generate() else sm(h))
+        h = self.norm2(x)
+        cm = self.channel_mixer
+        x = x + (cm.generate(h) if cm.has_custom_generate() else cm(h))
+        return x + self.x0_lambda * x0
+
+
+class PiCOFormer(VathosModel):
+    """High-level transformer. Ultra-efficient, modular, x0-skip nativo.
+
+    Caratteristiche:
+    - `d_model` fissa su tutti i layer.
+    - Spatial mixers heterogeneous (lista di `Builder`, uno per layer).
+    - x0-skip nativo via `BlockX0` / `PiCOBlock`.
+    - Value embeddings opzionali (zero cost se off, indexed-by-layer).
+    - Smear gate opzionale (Modded-NanoGPT style) su x0.
+    - Final logit softcap.
+    - Tied embeddings di default.
+
+    Args:
+        vocab_size: vocab size.
+        d_model: residual stream dim (fissa).
+        n_layers: numero blocchi.
+        spatials: `Builder` (condiviso) o `List[Builder]` di lunghezza n_layers.
+        channel: `Builder` channel mixer (default VariableUDLP M=4*d_model, ReLU²).
+        norm: classe norm (default RMSNorm).
+        ve_groups: `List[int|None]` di lunghezza n_layers. `None` = no VE su quel layer;
+                   int = id gruppo (gruppi condividono lo stesso `nn.Embedding`).
+                   `None` globalmente = niente VE (default, niente memory cost).
+        smear_gate: `nn.Module` o `None`. Applicato a x0 dopo l'embedding.
+        logit_softcap: softcap finale. 0 disabilita.
+        tied_embeddings: tying input/output embedding (default True).
+    """
+    __name__ = "PiCOFormer"
+
+    def __init__(self, vocab_size: int, d_model: int, n_layers: int,
+                 spatials, channel=None, norm=RMSNorm,
+                 ve_groups=None, smear_gate: nn.Module = None,
+                 smear_gate_lookback: int = 0,
+                 logit_softcap: float = 30.0, tied_embeddings: bool = True):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.softcap = logit_softcap
+
+        # Spatials: lista di Builder o Builder singolo (condiviso)
+        spatials_list = spatials if isinstance(spatials, list) else [spatials] * n_layers
+        assert len(spatials_list) == n_layers, "len(spatials) must equal n_layers"
+
+        # Channel default = VariableUDLP, M=4*d_model, ReLU²
+        if channel is None:
+            channel = Builder(VariableUDLP, d_output=d_model, M=4 * d_model, activation=ReLU2)
+
+        # VE: pre-compute embedding condiviso per group + indice per layer
+        self._has_any_ve = ve_groups is not None and any(g is not None for g in ve_groups)
+        if self._has_any_ve:
+            assert len(ve_groups) == n_layers, "len(ve_groups) must equal n_layers"
+            unique_groups = sorted({g for g in ve_groups if g is not None})
+            self.ve_embeddings = nn.ModuleList([nn.Embedding(vocab_size, d_model) for _ in unique_groups])
+            group_to_idx = {g: i for i, g in enumerate(unique_groups)}
+            self._ve_idx_per_layer = tuple(
+                group_to_idx[g] if g is not None else -1 for g in ve_groups
+            )
+        else:
+            self.ve_embeddings = nn.ModuleList()
+            self._ve_idx_per_layer = (-1,) * n_layers
+        self._ve_enabled_per_layer = tuple(idx >= 0 for idx in self._ve_idx_per_layer)
+
+        # Embeddings + final norm + unembedder
+        self.embedder = Embedder(vocab_size, d_model)
+        self.unembedder = UnbiasedLinear(d_model, vocab_size)
+        if tied_embeddings:
+            self.unembedder.linear.weight = self.embedder.embedding.weight
+        self.final_norm = norm(d_model)
+
+        # Blocks
+        self.blocks = nn.ModuleList([
+            PiCOBlock(
+                d_model,
+                channel_mixer=channel(d_model),
+                spatial_mixer=spatials_list[i](d_model),
+                norm=norm,
+                ve_enabled=self._ve_enabled_per_layer[i],
+            )
+            for i in range(n_layers)
+        ])
+
+        # Smear gate opzionale (nn.Module [B,L,D] -> [B,L,D]).
+        # smear_gate_lookback: 0 = stateless (per-token pointwise) — il fast path
+        # generate() lo applica al singolo token. Se >0, il fast path mantiene una
+        # rolling window di `lookback` token raw embeddings e ricomputa smear su
+        # `window + new_token` (necessario per conv causali kernel>1, RNN, ecc.).
+        self.smear_gate = smear_gate
+        self._smear_lookback = int(smear_gate_lookback)
+        self._x_window = None
+        self._smear_gate_warned = False
+
+        self._init_weights()
+
+    def _init_weights(self):
+        # Embeddings: normal std=0.02 (standard). LM head tied di default ⇒ stesso peso.
+        std_embed = 0.02
+        nn.init.normal_(self.embedder.embedding.weight, mean=0.0, std=std_embed)
+        if self.unembedder.linear.weight is not self.embedder.embedding.weight:
+            nn.init.normal_(self.unembedder.linear.weight, mean=0.0, std=std_embed)
+        for emb in self.ve_embeddings:
+            nn.init.normal_(emb.weight, mean=0.0, std=std_embed)
+
+        # Sub-modules: ciascuno definisce la propria policy di identity-init.
+        if self.smear_gate is not None and hasattr(self.smear_gate, "_init_weights"):
+            self.smear_gate._init_weights()
+        for block in self.blocks:
+            if hasattr(block.channel_mixer, "_init_weights"):
+                block.channel_mixer._init_weights()
+            if hasattr(block.spatial_mixer, "_init_weights"):
+                block.spatial_mixer._init_weights()
+        # x0_lambda (BlockX0) e ve_scale (PiCOBlock) sono già zero-init da nn.Parameter(torch.zeros(1)).
+
+    def _compute_ves(self, input_ids: torch.Tensor):
+        """Pre-computa la lista di tensori VE per ogni layer (None se non abilitato)."""
+        if not self._has_any_ve:
+            return (None,) * self.n_layers
+        # Computa unique embeddings UNA volta, poi distribuisce per layer
+        unique_ves = [emb(input_ids) for emb in self.ve_embeddings]
+        return tuple(unique_ves[idx] if idx >= 0 else None for idx in self._ve_idx_per_layer)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_ids = x
+        x0 = self.embedder(x)
+        if self.smear_gate is not None:
+            x0 = self.smear_gate(x0)
+
+        ves = self._compute_ves(input_ids)
+
+        x = x0
+        for i, block in enumerate(self.blocks):
+            x = block(x, x0, ve=ves[i])
+
+        logits = self.unembedder(self.final_norm(x))
+        if self.softcap > 0:
+            logits = self.softcap * torch.tanh(logits / self.softcap)
+        return logits
+
+    @torch.no_grad()
+    def _clear_caches(self):
+        for block in self.blocks:
+            for mod in block.modules():
+                if hasattr(mod, "clear_cache"):
+                    mod.clear_cache()
+        self._x_window = None
+
+    @torch.no_grad()
+    def generate(self, prompt: torch.Tensor, max_len: int = 100,
+                 temperature: float = 1.0, top_k: int | None = None,
+                 top_p: float = 1.0, repetition_penalty: float = 1.0,
+                 token_end=None, simple: bool = False) -> torch.Tensor:
+        """KV-cached AR generation. Prefill in un colpo, poi token-by-token."""
+        if simple:
+            return simple_ar_generate(
+                self, prompt, max_len=max_len, temperature=temperature,
+                top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty,
+                token_end=token_end,
+            )
+
+        self.eval()
+        self._clear_caches()
+
+        # Warning una tantum se smear_gate è non-trivial ma lookback=0
+        if (self.smear_gate is not None
+                and not isinstance(self.smear_gate, nn.Identity)
+                and self._smear_lookback == 0
+                and not self._smear_gate_warned):
+            import warnings
+            warnings.warn(
+                "PiCOFormer.generate(simple=False): smear_gate non-Identity con "
+                "smear_gate_lookback=0. Se il modulo ha state temporale (conv "
+                "kernel>1, RNN, ecc.), il fast path divergerà da simple=True. "
+                "Setta smear_gate_lookback=<receptive_field-1> o usa simple=True.",
+                RuntimeWarning, stacklevel=2,
+            )
+            self._smear_gate_warned = True
+
+        if prompt.dim() == 1:
+            prompt = prompt.unsqueeze(0)
+        device = next(self.parameters()).device
+        prompt = prompt.to(device)
+
+        # --- Prefill --------------------------------------------------------
+        x0_raw = self.embedder(prompt)
+        if self.smear_gate is not None:
+            x0 = self.smear_gate(x0_raw)
+            if self._smear_lookback > 0:
+                # Salva ultimi `lookback` raw embeddings come window (esclude i nuovi)
+                self._x_window = x0_raw[:, -self._smear_lookback:, :].clone()
+        else:
+            x0 = x0_raw
+        ves = self._compute_ves(prompt)
+        x = x0
+        for i, block in enumerate(self.blocks):
+            x = block.generate(x, x0, ve=ves[i])
+        logits = self.unembedder(self.final_norm(x))[:, -1, :]
+        if self.softcap > 0:
+            logits = self.softcap * torch.tanh(logits / self.softcap)
+        if repetition_penalty != 1.0:
+            logits = apply_repetition_penalty(logits, prompt, repetition_penalty)
+        next_token = sample_next_token(logits, temperature, top_k, top_p)
+        generated = torch.cat([prompt, next_token], dim=1)
+
+        # --- Step-by-step ---------------------------------------------------
+        pbar = tqdm(range(max_len - 1), desc="PiCOFormer fast gen")
+        for _ in pbar:
+            x0_t_raw = self.embedder(next_token)
+            if self.smear_gate is not None:
+                if self._smear_lookback > 0:
+                    # combina window + new, applica smear, prendi ultimo token
+                    combined = torch.cat([self._x_window, x0_t_raw], dim=1)
+                    smeared = self.smear_gate(combined)
+                    x0_t = smeared[:, -1:, :]
+                    # aggiorna window: tieni gli ultimi `lookback` raw embeddings
+                    self._x_window = combined[:, -self._smear_lookback:, :]
+                else:
+                    x0_t = self.smear_gate(x0_t_raw)
+            else:
+                x0_t = x0_t_raw
+            ves_t = self._compute_ves(next_token)
+            x_t = x0_t
+            for i, block in enumerate(self.blocks):
+                x_t = block.generate(x_t, x0_t, ve=ves_t[i])
+            logits = self.unembedder(self.final_norm(x_t))[:, -1, :]
+            if self.softcap > 0:
+                logits = self.softcap * torch.tanh(logits / self.softcap)
+            if repetition_penalty != 1.0:
+                logits = apply_repetition_penalty(logits, generated, repetition_penalty)
+            next_token = sample_next_token(logits, temperature, top_k, top_p)
+            generated = torch.cat([generated, next_token], dim=1)
+            if token_end is not None and (next_token == token_end).all():
+                break
+
+        self._clear_caches()
+        return generated
+
+
+if __name__ == "__main__":
+    d_models = [64 for i in range(4)]
+    m_dims = [128, 64, 32, 16]
+    attn = Builder(GQA, n_heads=4, n_kv_heads=2)
+    model = ModdedFormer(
+        vocab_size=100,
+        embed_dim=64,
+        d_models=d_models,
+        spatials=[attn for _ in range(len(d_models))],
+        M_dims=m_dims,
+        input_projections=3,
+        zeroskip=True
+    )
+    model.summary()
+
+    out = model(
+        torch.randint(0, 99, (2, 128))
+    )
+    print(out.shape)
+
+    model.profile()

vathos/complexity.py

@@ -0,0 +1,238 @@
+import re
+from collections import defaultdict
+from typing import List, Tuple, Dict
+
+
+def parse_complexity(complexity: str) -> Tuple[Dict[str, int], List[str]]:
+    """
+    Parse a Big O complexity string.
+
+    Returns:
+        Tuple of (variable_exponents dict, log_terms list)
+        where log_terms contains the variables inside log (e.g., ['n', 'L'])
+    """
+    match = re.search(r'O\((.*?)\)', complexity)
+    if not match:
+        return {}, []
+
+    content = match.group(1).strip()
+
+    log_terms = []
+    log_matches = re.findall(r'\blog\s+(\w+)', content, flags=re.IGNORECASE)
+    log_terms.extend(log_matches)
+
+    standalone_logs = re.findall(r'\blog(?!\s+\w)', content, flags=re.IGNORECASE)
+    log_terms.extend(['n'] * len(standalone_logs))
+
+    content_no_log = re.sub(r'\blog\s*\w*', '', content, flags=re.IGNORECASE)
+
+    exponents = defaultdict(int)
+    tokens = re.findall(r'([a-zA-Z_]\w*)(?:\^(\d+))?', content_no_log)
+
+    for var, exp in tokens:
+        if var and var.lower() != 'log':
+            exponent = int(exp) if exp else 1
+            exponents[var] = max(exponents[var], exponent)
+
+    return dict(exponents), log_terms
+
+
+def compare_complexities(c1: Tuple[Dict[str, int], List[str]],
+                         c2: Tuple[Dict[str, int], List[str]]) -> int:
+    """
+    Compare two complexities.
+    Returns: 1 if c1 > c2, -1 if c1 < c2, 0 if equal
+    """
+    exp1, log1 = c1
+    exp2, log2 = c2
+
+    all_vars = set(exp1.keys()) | set(exp2.keys())
+
+    for var in sorted(all_vars):
+        e1 = exp1.get(var, 0)
+        e2 = exp2.get(var, 0)
+        if e1 > e2:
+            return 1
+        elif e1 < e2:
+            return -1
+
+    if len(log1) > len(log2):
+        return 1
+    elif len(log1) < len(log2):
+        return -1
+
+    return 0
+
+
+def combine_big_o_product(complexities: List[str]) -> str:
+    """
+    Combines multiple Big O complexity strings by taking the dominant complexity.
+
+    Args:
+        complexities: List of Big O notation strings
+
+    Returns:
+        Combined Big O notation representing the dominant complexity
+
+    Examples:
+        >>> combine_big_o(["O(n^2)", "O(n log n)", "O(n^3)"])
+        'O(n^3)'
+        >>> combine_big_o(["O(n)", "O(n log n)"])
+        'O(n log n)'
+        >>> combine_big_o(["O(L^2 d^2)", "O(L d^2 k)", "O(L^2 d^2)"])
+        'O(L^2 d^2 k)'
+    """
+    if not complexities:
+        return "O(1)"
+
+    parsed = [parse_complexity(c) for c in complexities]
+
+    # Find the dominant complexity
+    dominant = parsed[0]
+    dominant_str = complexities[0]
+
+    for i in range(1, len(parsed)):
+        if compare_complexities(parsed[i], dominant) > 0:
+            dominant = parsed[i]
+            dominant_str = complexities[i]
+
+    # Find the maximum exponents for all variables
+    max_exponents = defaultdict(int)
+    for exponents, log_terms in parsed:
+        for var, exp in exponents.items():
+            max_exponents[var] = max(max_exponents[var], exp)
+
+    # Collect all log terms from complexities with maximum polynomial degree
+    all_log_terms = []
+    for exponents, log_terms in parsed:
+        has_max_degree = True
+
+        # Check if this has any variable at less than max exponent
+        for var, exp in exponents.items():
+            if exp < max_exponents[var]:
+                has_max_degree = False
+                break
+
+        # If it has max degree for its variables, keep its log terms
+        if has_max_degree:
+            all_log_terms.extend(log_terms)
+
+    if not max_exponents and not all_log_terms:
+        return "O(1)"
+
+    # Build the result string
+    terms = []
+
+    # Add variable terms
+    for var in sorted(max_exponents.keys()):
+        exp = max_exponents[var]
+        if exp == 1:
+            terms.append(var)
+        else:
+            terms.append(f"{var}^{exp}")
+
+    # Add log terms if applicable
+    if all_log_terms:
+        log_str = " ".join([f"log {var}" for var in sorted(set(all_log_terms))])
+        if terms:
+            terms.append(log_str)
+        else:
+            return f"O({log_str})"
+
+    return f"O({' '.join(terms)})"
+
+
+from typing import List
+from collections import defaultdict
+import re
+
+
+def combine_big_o_sum(complexities: List[str]) -> str:
+    """
+    Combines multiple Big O complexity strings by summing them and simplifying
+    like terms to return a human-readable dominant complexity.
+
+    Args:
+        complexities: List of Big O notation strings
+
+    Returns:
+        Combined Big O notation representing the simplified sum
+
+    Examples:
+        >>> combine_big_o_sum(["O(n^2)", "O(n)", "O(1)"])
+        'O(n^2)'
+        >>> combine_big_o_sum(["O(L d^2)", "O(L d^2)", "O(L^2 d)"])
+        'O(L d^2 + L^2 d)'
+        >>> combine_big_o_sum(["O(depth L d^2)", "O(L^2 d)", "O(L d^2)", "O(L d^2)", "O(L d)"])
+        'O(L d^2 + L^2 d)'
+    """
+    if not complexities:
+        return "O(1)"
+
+    all_terms = []
+    for c in complexities:
+        term = c.strip()
+        if term.startswith("O(") and term.endswith(")"):
+            term = term[2:-1].strip()
+        if term and term != "1":
+            # Split by + to handle multiple terms in one complexity
+            all_terms.extend([t.strip() for t in term.split('+')])
+
+    if not all_terms:
+        return "O(1)"
+
+    term_groups = defaultdict(list)
+
+    for term in all_terms:
+        var_signature = []
+        parts = term.split()
+
+        for part in parts:
+            if part.replace('.', '').isdigit() or part in ['depth', 'batch', 'const']:
+                continue
+
+            # Parse variable with optional exponent
+            if '^' in part:
+                var, exp = part.split('^')
+                var_signature.append((var, int(exp)))
+            else:
+                var_signature.append((part, 1))
+
+        # Sort to create canonical signature
+        var_signature = tuple(sorted(var_signature))
+        term_groups[var_signature].append(term)
+
+    unique_terms = []
+    for sig, terms in term_groups.items():
+        representative = min(terms, key=lambda t: (len(t), t))
+        unique_terms.append((sig, representative))
+
+    if not unique_terms:
+        return "O(1)"
+
+    def complexity_score(sig_term):
+        sig, term = sig_term
+        if not sig:
+            return (0, 0, term)
+        # Total degree and max degree
+        total_degree = sum(exp for var, exp in sig)
+        max_degree = max(exp for var, exp in sig)
+        return (max_degree, total_degree, term)
+
+    unique_terms.sort(key=complexity_score, reverse=True)
+
+    # If clearly dominated, return just the dominant term(s)
+    if len(unique_terms) > 1:
+        top_score = complexity_score(unique_terms[0])
+        # Keep all terms with the same max complexity
+        result_terms = [term for sig_term in unique_terms
+                        if complexity_score(sig_term)[0] == top_score[0]]
+    else:
+        result_terms = [unique_terms[0][1]]
+
+    # Build result
+    if len(result_terms) == 1:
+        return f"O({result_terms[0]})"
+    else:
+        return f"O({' + '.join(result_terms)})"
+

vathos/convolutions.py

@@ -0,0 +1,119 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class FastShiftConv1d_K3(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.w0 = nn.Parameter(torch.ones(dim))
+        self.w1 = nn.Parameter(torch.zeros(dim))
+        self.w2 = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1 = F.pad(x[:, :-1, :], (0, 0, 1, 0))
+        x2 = F.pad(x[:, :-2, :], (0, 0, 2, 0))
+
+        return (x * self.w0.to(x.dtype)) + (x1 * self.w1.to(x.dtype)) + (x2 * self.w2.to(x.dtype))
+
+
+class FastShiftConv1d_K4(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.w0 = nn.Parameter(torch.ones(dim))
+        self.w1 = nn.Parameter(torch.zeros(dim))
+        self.w2 = nn.Parameter(torch.zeros(dim))
+        self.w3 = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1 = F.pad(x[:, :-1, :], (0, 0, 1, 0))
+        x2 = F.pad(x[:, :-2, :], (0, 0, 2, 0))
+        x3 = F.pad(x[:, :-3, :], (0, 0, 3, 0))
+
+        return (x * self.w0.to(x.dtype)) + \
+            (x1 * self.w1.to(x.dtype)) + \
+            (x2 * self.w2.to(x.dtype)) + \
+            (x3 * self.w3.to(x.dtype))
+
+
+class FastShiftConv1d_K2(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        # Dirac/Identity Init: 1 for current token, 0 for past tokens
+        self.w0 = nn.Parameter(torch.ones(dim))
+        self.w1 = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1 = F.pad(x[:, :-1, :], (0, 0, 1, 0))
+
+        return (x * self.w0.to(x.dtype)) + (x1 * self.w1.to(x.dtype))
+
+
+class RWKVTimeMix(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.time_mix = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_prev = F.pad(x[:, :-1, :], (0, 0, 1, 0))
+        mix = self.time_mix.to(x.dtype)
+
+        return x * mix + x_prev * (1.0 - mix)
+
+
+class SmearGate2(nn.Module):
+    def __init__(self, dim: int, gate_dim: int = 12):
+        super().__init__()
+        self.gate_dim = min(gate_dim, dim)
+        self.W_g = nn.Parameter(torch.zeros(self.gate_dim, dim))
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_prev = F.pad(x[:, :-1, :], (0, 0, 1, 0))
+        x_slice = x[..., :self.gate_dim]
+        g = torch.sigmoid(torch.matmul(x_slice, self.W_g.to(x.dtype)))
+
+        return x + self.smear_lambda * g * x_prev
+
+
+class SmearGateLookback3(nn.Module):
+    def __init__(self, dim: int, gate_dim: int = 12):
+        super().__init__()
+        self.gate_dim = min(gate_dim, dim)
+
+        self.W_g = nn.Parameter(torch.zeros(self.gate_dim, dim * 3))
+
+        self.smear_lambdas = nn.Parameter(torch.zeros(3))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_slice = x[..., :self.gate_dim]
+        g_all = torch.sigmoid(torch.matmul(x_slice, self.W_g.to(x.dtype)))
+
+        g1, g2, g3 = g_all.chunk(3, dim=-1)
+
+        x1 = F.pad(x[:, :-1, :], (0, 0, 1, 0))
+        x2 = F.pad(x[:, :-2, :], (0, 0, 2, 0))
+        x3 = F.pad(x[:, :-3, :], (0, 0, 3, 0))
+
+        return x + (self.smear_lambdas[0] * g1 * x1) \
+            + (self.smear_lambdas[1] * g2 * x2) \
+            + (self.smear_lambdas[2] * g3 * x3)
+
+
+class ModdedSmearGate(nn.Module):
+    def __init__(self, gate_dim: int = 12):
+        super().__init__()
+        self.gate_dim = gate_dim
+
+        self.smear_gate = nn.Linear(gate_dim, 1, bias=False)
+
+        with torch.no_grad():
+            nn.init.zeros_(self.smear_gate.weight)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_prev = F.pad(x[:, :-1, :], (0, 0, 1, 0))
+        x_slice = x[..., :self.gate_dim]
+
+        g = torch.sigmoid(self.smear_gate(x_slice))
+        return x + g * (x_prev - x)

vathos/data.py

@@ -0,0 +1,170 @@
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoTokenizer
+from dataclasses import dataclass
+from typing import List, Dict, Any, Union, Optional
+from pathlib import Path
+
+
+class TextDataset(Dataset):
+    """Text dataset supporting lists, files, and folder structures."""
+
+    def __init__(self, source: Union[List[str], str, Path]):
+        """
+        Args:
+            source: Can be:
+                - List of strings (text data)
+                - Path to single .txt file
+                - Path to folder containing .txt files (recursive)
+        """
+        self.texts = []
+        self.file_paths = []  # Track source files for debugging
+
+        if isinstance(source, list):
+            self.texts = source
+            self.file_paths = [None] * len(source)
+        else:
+            source = Path(source)
+            if source.is_file():
+                self._load_file(source)
+            elif source.is_dir():
+                self._load_directory(source)
+            else:
+                raise ValueError(f"Source not found: {source}")
+
+        if len(self.texts) == 0:
+            raise ValueError("No text data loaded!")
+
+    def _load_file(self, file_path: Path):
+        """Load a single text file."""
+        with open(file_path, 'r', encoding='utf-8') as f:
+            text = f.read().strip()
+            if text:
+                self.texts.append(text)
+                self.file_paths.append(str(file_path))
+
+    def _load_directory(self, dir_path: Path):
+        """Recursively load all .txt files from directory."""
+        txt_files = sorted(dir_path.rglob("*.txt"))
+
+        for file_path in txt_files:
+            try:
+                self._load_file(file_path)
+            except Exception as e:
+                print(f"Warning: Failed to load {file_path}: {e}")
+
+    def __len__(self):
+        return len(self.texts)
+
+    def __getitem__(self, idx):
+        return self.texts[idx]
+
+    def get_file_path(self, idx: int) -> Optional[str]:
+        """Get source file path for debugging."""
+        return self.file_paths[idx]
+
+
+@dataclass
+class TokenizerConfig:
+    """Tokenizer configuration."""
+    model_name: str = "bert-base-uncased"
+    max_length: int = 512
+    truncation: bool = True
+    padding: str = "max_length"  # or "longest"
+    return_tensors: str = "pt"
+
+
+class TokenizerWrapper:
+    """Wrapper for HuggingFace tokenizer with collating dataloader."""
+
+    def __init__(self, config: TokenizerConfig = None):
+        self.config = config or TokenizerConfig()
+        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_name)
+
+    def collate_fn(self, batch: List[str]) -> Dict[str, Any]:
+        """Collate function for DataLoader."""
+        return self.tokenizer(
+            batch,
+            max_length=self.config.max_length,
+            truncation=self.config.truncation,
+            padding=self.config.padding,
+            return_tensors=self.config.return_tensors
+        )
+
+    def decode(self, token_ids, skip_special_tokens: bool = True) -> Union[str, List[str]]:
+        """
+        Decode token IDs back to text.
+
+        Args:
+            token_ids: Single sequence or batch of token IDs (tensor or list)
+            skip_special_tokens: Whether to remove [CLS], [SEP], [PAD] etc.
+
+        Returns:
+            Decoded text string or list of strings
+        """
+        return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
+
+    def batch_decode(self, token_ids, skip_special_tokens: bool = True) -> List[str]:
+        """
+        Decode a batch of token IDs back to texts.
+
+        Args:
+            token_ids: Batch of token IDs (shape: [batch_size, seq_len])
+            skip_special_tokens: Whether to remove special tokens
+
+        Returns:
+            List of decoded text strings
+        """
+        return self.tokenizer.batch_decode(token_ids, skip_special_tokens=skip_special_tokens)
+
+    def get_dataloader(
+        self,
+        source: Union[List[str], str, Path],
+        batch_size: int = 8,
+        shuffle: bool = True,
+        num_workers: int = 0
+    ) -> DataLoader:
+        """
+        Create a DataLoader with tokenization collation.
+
+        Args:
+            source: List of texts, file path, or directory path
+            batch_size: Batch size
+            shuffle: Whether to shuffle data
+            num_workers: Number of workers for data loading
+        """
+        dataset = TextDataset(source)
+        print(f"Loaded {len(dataset)} texts")
+
+        return DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            num_workers=num_workers,
+            collate_fn=self.collate_fn
+        )
+
+
+if __name__ == "__main__":
+    texts = "_tests"
+
+    config = TokenizerConfig(
+        model_name="bert-base-uncased",
+        max_length=32,
+        padding="longest"
+    )
+    wrapper = TokenizerWrapper(config)
+
+    print("=== Example 1: List of texts ===")
+    dataloader = wrapper.get_dataloader(texts, batch_size=2, shuffle=False)
+
+    for batch_idx, batch in enumerate(dataloader):
+        print(f"\nBatch {batch_idx}: {batch['input_ids'].shape}")
+
+        decoded_batch = wrapper.batch_decode(batch['input_ids'], skip_special_tokens=True)
+        print(f"Decoded texts: {decoded_batch}")
+
+        single_decoded = wrapper.decode(batch['input_ids'][0], skip_special_tokens=True)
+        print(f"First sequence: {single_decoded}")
+
+        with_special = wrapper.decode(batch['input_ids'][0], skip_special_tokens=False)
+        print(f"With special tokens: {with_special}")

vathos/functions.py

@@ -0,0 +1,586 @@
+import torch
+import torch.nn.functional as F
+import time
+import numpy as np
+
+# matplotlib è usato solo dalle utility di plot/benchmark, non dal codice del modello.
+# Lazy + optional import per non rompere l'import di Vathos in ambienti headless.
+try:
+    import matplotlib.pyplot as plt
+    plt.style.use('ggplot')
+except ImportError:
+    plt = None  # le funzioni di plot solleveranno errore esplicito se chiamate
+
+try:
+    from colorama import Fore
+
+    GOOD = Fore.GREEN
+    BAD = Fore.RED
+    RES = Fore.RESET
+    SEC = Fore.LIGHTBLACK_EX
+    HM = Fore.YELLOW
+    NUM = Fore.BLUE
+except:
+    GOOD = ""
+    BAD = ""
+    RES = ""
+    SEC = ""
+    NUM = ""
+    HM = ""
+
+FLAG_PASS = 3
+
+
+def _require_matplotlib():
+    if plt is None:
+        raise ImportError(
+            "Funzione di plotting chiamata ma matplotlib non è installato. "
+            "Installa con `pip install matplotlib` per usarla."
+        )
+
+
+def plot(*args, **kwargs):
+    _require_matplotlib()
+    for x_i in args:
+        plt.plot(x_i)
+    for x_i in kwargs:
+        plt.plot(kwargs[x_i], label=str(x_i))
+
+    plt.legend()
+    plt.show()
+
+
+def cumsum(x):
+    return torch.cumsum(x, dim=1)
+
+
+def batched_channelwise_conv1d(K, V):
+    """
+    Computes a convolution between two batched tensors
+    K: [B, L, d] kernel
+    V: [B, L, d] signal
+    Output: [B, L, d]
+    """
+    B, L, d = V.shape
+    V_reshaped = V.view(1, B * L, d)
+
+    K_reshaped = K.view(B * L, 1, d).flip(-1)
+
+    pad_total = d - 1
+    pad_left = pad_total // 2
+    pad_right = pad_total - pad_left
+    V_padded = F.pad(V_reshaped, (pad_left, pad_right))
+
+    output = F.conv1d(V_padded, K_reshaped, groups=B * L)
+
+    return output.view(B, L, d)
+
+
+def holographic_binding(K, V):
+    """
+    Computes Batched Circular Convolution (Binding) using FFT.
+    This is O(d log d) and numerically stable.
+
+    K: [..., d] kernel
+    V: [..., d] signal
+    Output: [..., d]
+    """
+    V_f = torch.fft.rfft(V, dim=-1, norm='ortho')
+    K_f = torch.fft.rfft(K, dim=-1, norm='ortho')
+
+    output_f = V_f * K_f
+    output = torch.fft.irfft(output_f, n=V.shape[-1], dim=-1, norm='ortho')
+
+    return output
+
+
+def power_weigthed_cumsum(x, a=0.999, rescale=True):
+    """
+    a torch implementation of a power weighted cumsum,
+    which can be applied to batched inputs of shapes [B, L, d] or [B, L, d1, d2], summing over L
+    """
+    if a - 1 == 0:
+        return cumsum(x)
+    if len(x.shape) == 3:
+        alpha_pow = torch.full([x.shape[1]], a, dtype=x.dtype, device=x.device).cumprod(dim=0).unsqueeze(0).unsqueeze(2)
+    elif len(x.shape) == 4:
+        alpha_pow = torch.full([x.shape[1]], a, dtype=x.dtype, device=x.device).cumprod(dim=0).unsqueeze(0).unsqueeze(
+            2).unsqueeze(2)
+    else:
+        raise RuntimeError("x shapes mus be of form [B, L, d] or [B, L, d, d]")
+
+    return torch.cumsum(x / alpha_pow, dim=1) * alpha_pow * (alpha_pow[:, 0, ...] if rescale else 1)
+
+
+def stable_power_weighted_cumsum(x: torch.Tensor, a: float = 0.999, rescale: bool = True) -> torch.Tensor:
+    B, L = x.shape[:2]
+
+    dtype = x.dtype
+    if L > 1024 and x.dtype == torch.float32:
+        x = x.to(torch.float64)
+        a_tensor = torch.tensor(a, dtype=torch.float64, device=x.device)
+    else:
+        a_tensor = torch.tensor(a, dtype=x.dtype, device=x.device)
+
+    t = torch.arange(1, L + 1, dtype=a_tensor.dtype, device=x.device)
+
+    alpha_pow = torch.pow(a_tensor, t)
+
+    reshape_dims = [1, L] + [1] * (len(x.shape) - 2)
+    alpha_pow = alpha_pow.view(*reshape_dims)
+
+    scaled_input = x / alpha_pow
+
+    s_cumulative = torch.cumsum(scaled_input, dim=1)
+
+    s = s_cumulative * alpha_pow
+
+    if rescale:
+        s = s * (1.0 - a_tensor)
+
+    return s.to(dtype)
+
+
+def precompute_power_weigthed_cumsum(x, alpha_pow):
+    return torch.cumsum(x / alpha_pow, dim=1) * alpha_pow
+
+
+def flag(text, level=1):
+    if level <= FLAG_PASS:
+        print(f"{SEC}||{HM}FLAG LV.{level}{SEC}||{HM} {text}{RES}")
+
+
+def getname(obj):
+    return obj.__name__ if hasattr(obj, "__name__") else str(type(obj)).split('.')[-1]
+
+
+def benchmark_ar_symbolic_model(model, n, input_shape):
+    """
+    Benchmarks the autoregressive symbolic model on three tasks: forward pass,
+    forward + backward, and forward + backward + optimizer step. Uses random
+    generated data. Reports average times and shows a bar plot.
+
+    Additionally, benchmarks the ratio of forward time / (forward + backward) time
+    with increasing sequence length L, and plots it.
+
+    Args:
+        model: The model to benchmark (e.g., SequenceModel instance).
+        n: Number of iterations to average over for each benchmark.
+        input_shape: Tuple (batch_size, seq_len) for initial input shape.
+    """
+    device = next(model.parameters()).device
+    is_cuda = device.type == 'cuda'
+    model.train()  # Ensure model is in training mode for gradients
+    optimizer = torch.optim.Adam(model.parameters())
+
+    # Helper to sync if CUDA
+    def sync():
+        if is_cuda:
+            torch.cuda.synchronize()
+
+    # Part 1: Benchmark forward, forward+backward, forward+backward+opt
+    times_fwd = []
+    times_fb = []
+    times_fbo = []
+
+    for _ in range(n):
+        x = torch.randint(0, model.vocab_size, input_shape, device=device, requires_grad=False)
+
+        # Forward only
+        sync()
+        start = time.perf_counter()
+        out = model(x)
+        sync()
+        end = time.perf_counter()
+        times_fwd.append(end - start)
+
+        # Forward + backward
+        optimizer.zero_grad()
+        sync()
+        start = time.perf_counter()
+        out = model(x)
+        loss = out.sum()  # Dummy loss
+        loss.backward()
+        sync()
+        end = time.perf_counter()
+        times_fb.append(end - start)
+
+        # Forward + backward + optimizer step
+        optimizer.zero_grad()
+        sync()
+        start = time.perf_counter()
+        out = model(x)
+        loss = out.sum()
+        loss.backward()
+        optimizer.step()
+        sync()
+        end = time.perf_counter()
+        times_fbo.append(end - start)
+
+    avg_fwd = np.mean(times_fwd)
+    avg_fb = np.mean(times_fb)
+    avg_fbo = np.mean(times_fbo)
+
+    print(f"Average Forward Time: {avg_fwd:.6f} s")
+    print(f"Average Forward + Backward Time: {avg_fb:.6f} s")
+    print(f"Average Forward + Backward + Opt Time: {avg_fbo:.6f} s")
+
+    # Bar plot for the three tasks
+    labels = ['Forward', 'Forward + Backward', 'Forward + Backward + Opt']
+    times = [avg_fwd, avg_fb, avg_fbo]
+    plt.figure(figsize=(8, 5))
+    plt.bar(labels, times, color=['blue', 'orange', 'green'])
+    plt.ylabel('Average Time (s)')
+    plt.title(f'Benchmark Averages over {n} Iterations')
+    plt.show()
+
+    # Part 2: Ratio of forward / (forward + backward) with increasing L
+    batch_size = input_shape[0]
+    Ls = [16, 32, 64, 128, 256, 512, 1024]  # Reasonable sequence lengths, adjust as needed
+    ratios = []
+    avg_times_fwd = []
+    avg_times_fb = []
+    inner_n = max(10, n // 10)  # Average over fewer iterations for speed, but at least 10
+
+    for L in Ls:
+        shape = (batch_size, L)
+        x = torch.randint(0, model.vocab_size, shape, device=device, requires_grad=False)
+
+        # Time forward (average over inner_n)
+        fwd_times = []
+        for _ in range(inner_n):
+            sync()
+            start = time.perf_counter()
+            out = model(x)
+            sync()
+            end = time.perf_counter()
+            fwd_times.append(end - start)
+        time_fwd = np.mean(fwd_times)
+        avg_times_fwd.append(time_fwd)
+
+        # Time forward + backward (average over inner_n)
+        fb_times = []
+        for _ in range(inner_n):
+            optimizer.zero_grad()
+            sync()
+            start = time.perf_counter()
+            out = model(x)
+            loss = out.sum()
+            loss.backward()
+            sync()
+            end = time.perf_counter()
+            fb_times.append(end - start)
+        time_fb = np.mean(fb_times)
+        avg_times_fb.append(time_fb)
+
+        ratio = time_fwd / time_fb if time_fb > 0 else 0
+        ratios.append(ratio)
+
+    # Plot the ratio
+    plt.figure(figsize=(8, 5))
+    plt.plot(Ls, ratios, marker='o', color='red')
+    plt.xlabel('Sequence Length L')
+    plt.ylabel('Time Forward / Time (Forward + Backward)')
+    plt.title('Efficiency Ratio vs Sequence Length')
+    plt.grid(True)
+    plt.show()
+
+    # Optional: Plot absolute times for reference
+    plt.figure(figsize=(8, 5))
+    plt.plot(Ls, avg_times_fwd, marker='o', label='Forward', color='blue')
+    plt.plot(Ls, avg_times_fb, marker='o', label='Forward + Backward', color='orange')
+    plt.xlabel('Sequence Length L')
+    plt.ylabel('Average Time (s)')
+    plt.title('Absolute Times vs Sequence Length')
+    plt.legend()
+    plt.grid(True)
+    plt.show()
+
+
+def apply_repetition_penalty(logits: torch.Tensor, generated: torch.Tensor,
+                             penalty: float) -> torch.Tensor:
+    """
+    Standard HF-style repetition penalty, vectorized.
+    logits:    [B, V]   (last-position logits)
+    generated: [B, L]   (token ids generated so far, including prompt)
+    Divides logits>0 by penalty, multiplies logits<=0 by penalty, in-place-safe.
+    """
+    if penalty == 1.0:
+        return logits
+    score = torch.gather(logits, 1, generated)
+    score = torch.where(score > 0, score / penalty, score * penalty)
+    logits = logits.scatter(1, generated, score)
+    return logits
+
+
+def sample_next_token(logits: torch.Tensor, temperature: float = 1.0,
+                      top_k: int | None = None, top_p: float = 1.0) -> torch.Tensor:
+    """
+    Sample one token per row from [B, V] logits with temperature, top-k, top-p.
+    Returns [B, 1] long tensor.
+    """
+    logits = logits / max(temperature, 1e-8)
+
+    if top_k is not None and top_k > 0:
+        k = min(top_k, logits.size(-1))
+        v, _ = torch.topk(logits, k)
+        logits = logits.masked_fill(logits < v[:, [-1]], float('-inf'))
+
+    if top_p < 1.0:
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True)
+        cum = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        remove = cum > top_p
+        remove[..., 1:] = remove[..., :-1].clone()
+        remove[..., 0] = 0
+        mask = remove.scatter(1, sorted_idx, remove)
+        logits = logits.masked_fill(mask, float('-inf'))
+
+    probs = F.softmax(logits, dim=-1)
+    return torch.multinomial(probs, num_samples=1)
+
+
+@torch.no_grad()
+def simple_ar_generate(model, prompt, max_len=100, temperature=1.0,
+                       top_k=None, top_p=1.0, repetition_penalty=1.0,
+                       token_end=None, verbose=False):
+    """
+    Dumb autoregressive generation for debugging.
+
+    Runs a full forward pass on the whole sequence each step, takes logits at the
+    last position, samples, appends. O(L^2) total — useless for big models, but it
+    works on *anything* whose forward(ids) returns [B, L, V] logits, with no
+    assumptions about KV caches, custom generate(), pos encoders, ve, skips, etc.
+    Good as a ground-truth reference when debugging more sophisticated generators.
+    """
+    flag("simple_ar_generate is O(L^2) and ignores KV caches: use it for debugging, "
+         "not for real inference. If your model has a fast .generate(), prefer that.", 2)
+    model.eval()
+    if prompt.dim() == 1:
+        prompt = prompt.unsqueeze(0)
+    device = next(model.parameters()).device
+    generated = prompt.clone().to(device)
+
+    for _ in range(max_len):
+        logits = model(generated)
+        next_logits = logits[:, -1, :]
+
+        if repetition_penalty != 1.0:
+            next_logits = apply_repetition_penalty(next_logits, generated, repetition_penalty)
+
+        next_token = sample_next_token(next_logits, temperature, top_k, top_p)
+        generated = torch.cat([generated, next_token], dim=1)
+
+        if verbose:
+            print(next_token.flatten().tolist())
+
+        if token_end is not None and (next_token == token_end).all():
+            break
+
+    return generated
+
+
+def toeplitz_init(tensor: torch.Tensor, alpha: float, causal: bool = True, mul=0.1):
+    """
+    Initializes a square 2D tensor to a custom Toeplitz matrix with decaying kernels
+    generated from alpha^t for t=1...L, where L is the side length of the matrix.
+
+    If causal is True, the matrix is lower triangular (tril), containing values only
+    where row index i >= column index j.
+
+    Args:
+        tensor (torch.Tensor): The square 2D tensor to initialize in place.
+        alpha (float): The base for the exponential decay.
+        causal (bool, optional): If True, makes the matrix lower triangular. Defaults to False.
+
+    Returns:
+        torch.Tensor: The initialized tensor (modified in place).
+    """
+    if tensor.dim() != 2 or tensor.size(0) != tensor.size(1):
+        raise ValueError("Tensor must be a square 2D tensor.")
+
+    L = tensor.size(0)
+    i, j = torch.meshgrid(
+        torch.arange(L, device=tensor.device, dtype=tensor.dtype),
+        torch.arange(L, device=tensor.device, dtype=tensor.dtype),
+        indexing='ij'
+    )
+    dist = torch.abs(i - j)
+    t = dist + 1
+    matrix = alpha ** t
+
+    if causal:
+        matrix = torch.where(i >= j, matrix, torch.tensor(0.0, dtype=tensor.dtype, device=tensor.device))
+
+    with torch.no_grad():
+        tensor.copy_(matrix)
+
+    return tensor * mul
+
+
+def zero_layer_mtp_loss( # TODO: Vibecoded for now
+        logits: torch.Tensor,
+        targets: torch.Tensor,
+        mtp_weights: list[float],
+        ignore_index: int = -100,
+        softcap_val: float = 30.0
+) -> torch.Tensor:
+    """
+    Production-grade 0-layer MTP loss.
+
+    Optimizations:
+    1. Computes LogSumExp exactly once per token.
+    2. Uses index gathering for target logits to minimize memory bandwidth.
+    3. Casts to FP32 for numerical stability before reduction.
+    """
+    assert sum(mtp_weights) == 1, "the sum of MTP weights must be 1.0"
+
+    if softcap_val > 0.0:
+        logits = softcap_val * torch.tanh(logits / softcap_val)
+
+    lse = torch.logsumexp(logits, dim=-1).float()
+
+    total_loss = torch.tensor(0.0, device=logits.device, dtype=torch.float32)
+
+    for k, w in enumerate(mtp_weights):
+        if w <= 0.0:
+            continue
+
+        if k == 0:
+            valid_logits = logits
+            valid_targets = targets
+            valid_lse = lse
+        else:
+            valid_logits = logits[:, :-k]
+            valid_targets = targets[:, k:]
+            valid_lse = lse[:, :-k]
+
+        mask = (valid_targets != ignore_index)
+
+        safe_targets = torch.where(mask, valid_targets, torch.zeros_like(valid_targets))
+
+        target_logits = valid_logits.gather(
+            dim=-1,
+            index=safe_targets.unsqueeze(-1)
+        ).squeeze(-1).float()
+
+        ce_loss = valid_lse - target_logits
+
+        masked_loss = ce_loss * mask
+        num_valid_tokens = mask.sum().clamp(min=1)
+        loss_k = masked_loss.sum() / num_valid_tokens
+
+        total_loss += w * loss_k
+
+    return total_loss
+
+
+def plot_tensor_diagnostics(
+        name: str,
+        tensor: torch.Tensor,
+        on_grad: bool = False,
+        top_k_svd: int | None = None,
+        bins: int = 100
+):
+
+    has_grad = on_grad and hasattr(tensor, 'grad') and tensor.grad is not None
+    rows = 2 if has_grad else 1
+    fig, axes = plt.subplots(rows, 2, figsize=(12, 5 * rows))
+    if rows == 1:
+        axes = np.expand_dims(axes, axis=0)  # standardize indexing
+
+    def _compute_and_plot(t: torch.Tensor, row: int, title_prefix: str):
+        t_cpu = t.detach().float().cpu()
+
+        if t_cpu.ndim == 1:
+            mat = t_cpu.unsqueeze(0)
+        else:
+            mat = t_cpu.view(t_cpu.size(0), -1)
+
+        s_vals = torch.linalg.svdvals(mat).numpy()
+        if top_k_svd is not None:
+            s_vals = s_vals[:top_k_svd]
+
+        ax_svd = axes[row, 0]
+        ax_svd.plot(s_vals, marker='.', linestyle='-', color='b', markersize=4)
+        ax_svd.set_title(f"{title_prefix} Singular Values\nMax: {s_vals[0]:.4f}, Min: {s_vals[-1]:.4f}")
+        ax_svd.set_xlabel("Index")
+        ax_svd.set_ylabel("Singular Value $\sigma_i$")
+        ax_svd.set_yscale('log')  # Log scale is crucial to see the long tail
+        ax_svd.grid(True, which="both", ls="--", alpha=0.5)
+
+        ax_hist = axes[row, 1]
+        vals_flat = t_cpu.numpy().flatten()
+        ax_hist.hist(vals_flat, bins=bins, color='g', alpha=0.7, log=True)
+        mean, std = vals_flat.mean(), vals_flat.std()
+        ax_hist.set_title(f"{title_prefix} Distribution\n$\mu$: {mean:.2e}, $\sigma$: {std:.2e}")
+        ax_hist.set_xlabel("Value")
+        ax_hist.set_ylabel("Count (Log Scale)")
+        ax_hist.grid(True, ls="--", alpha=0.5)
+
+    _compute_and_plot(tensor, row=0, title_prefix=f"Weight: {name}")
+
+    if has_grad:
+        _compute_and_plot(tensor.grad, row=1, title_prefix=f"Gradient: {name}")
+    elif on_grad:
+        print(f"Warning: on_grad=True for '{name}', but tensor.grad is None.")
+
+    plt.tight_layout()
+    plt.show()
+
+
+import math
+
+def compute_normalized_grad_norm(model: torch.nn.Module) -> float:
+    """
+    Returns the RMS gradient norm across all trainable parameters.
+
+    RMS = sqrt( sum(g_i^2) / N )   where N = total number of scalar gradients.
+
+    Because this *averages* rather than sums, the result is independent of
+    model size: a 1 M-param and a 100 M-param model both live in [0, ~1]
+    during healthy training, making cross-run comparisons meaningful.
+
+    Returns 0.0 if no gradients are present yet.
+    """
+    total_sq = 0.0
+    total_n  = 0
+
+    for p in model.parameters():
+        if p.grad is not None:
+            total_sq += p.grad.detach().float().pow(2).sum().item()
+            total_n  += p.grad.numel()
+
+    if total_n == 0:
+        return 0.0
+
+    return math.sqrt(total_sq / total_n)
+
+
+def compute_layer_grad_norms(model: torch.nn.Module) -> dict[str, float]:
+    """
+    Returns a dict { layer_name -> RMS grad norm } for every named module
+    that owns at least one parameter with a gradient.
+
+    Same RMS normalisation as compute_normalized_grad_norm so per-layer
+    values are still comparable across models of different sizes.
+    Leaf modules only (skips container modules to avoid double-counting).
+    """
+    norms = {}
+
+    for name, module in model.named_modules():
+        # only leaf modules that directly own parameters
+        own_params = list(module.parameters(recurse=False))
+        if not own_params:
+            continue
+
+        sq_sum = 0.0
+        n      = 0
+        for p in own_params:
+            if p.grad is not None:
+                sq_sum += p.grad.detach().float().pow(2).sum().item()
+                n      += p.grad.numel()
+
+        if n > 0:
+            norms[name] = math.sqrt(sq_sum / n)
+
+    return norms

vathos/mamba.py

@@ -0,0 +1,27 @@
+from Vathos.blocks import *
+
+try:
+    from mamba_ssm import Mamba
+except ImportError:
+    raise ImportError("If you're trying to use Mamba, please install mamba-ssm, via pip install mamba-ssm. \n"
+                      "probably you'll need to also install the right torch version:\n"
+                      "pip install -q -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126")
+
+
+class MambaMixer(Layer):
+    def __init__(self, d_model, expand=4, d_state=64, k=3):
+        super(MambaMixer).__init__()
+        self.expand = expand
+        self.d_model = d_model
+        self.d_state = d_state
+        self.k = k
+        self.mamba = Mamba(
+            d_model=d_model,
+            d_state=d_state,
+            d_conv=d_state,
+            expand=expand,
+        )
+
+    def forward(self, x):
+        x = self.mamba(x)
+        return x

vathos/optimizers.py

@@ -0,0 +1,518 @@
+import random
+import torch
+import torch.optim as optim
+from Vathos.functions import flag
+from torch.optim import Optimizer, AdamW, SGD
+
+try:
+    from muon import Muon
+except ImportError:
+    flag("Muon library is not installed, usage of Muon dependent optimizer will not be possible", 2)
+
+
+class ValueScheduler:
+    def __init__(self):
+        self.step = 0
+        self.value = 0
+        self.functions = {}
+        self.values = []
+        self.compiled = False
+
+    def set(self, f, l=0, u=float('inf')):
+        assert l < u
+        self.functions[(l, u)] = f
+
+    def compile(self, l, m):
+        if self.compiled:
+            flag("ValueScheduler Already Compiled, recompiling...")
+        self.values = []
+        for l, u in self.functions:
+            i = l
+            f = self.functions[(l, u)]
+            while i <= u:
+                self.values.append(f(i))
+
+    def get(self, step):
+        assert self.compiled
+        return self.values[step]
+
+
+class SignSGD(optim.Optimizer):
+
+    def __init__(self, params, lr=0.01, rand_zero=True):
+        defaults = dict(lr=lr, rand_zero=rand_zero)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group['lr']
+            rand_zero = group['rand_zero']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = torch.sign(p.grad)
+
+                if rand_zero:
+                    zero_mask = (grad == 0)
+                    if zero_mask.any():
+                        grad[zero_mask] = torch.randint(
+                            0, 2, (zero_mask.sum().item(),),
+                            dtype=grad.dtype,
+                            device=grad.device
+                        ) * 2 - 1
+
+                p.add_(grad, alpha=-lr)
+
+        return loss
+
+
+"""
+stochastic_sign_muon.py
+=======================
+StochasticSignMuon: per-step stochastic interpolation between Muon and SignSGD
+for 2-D hidden weight matrices, with AdamW for everything else.
+
+Design principle
+----------------
+Both Muon and SignSGD operate on the *same* Nesterov momentum buffer:
+
+    buf  ←  β·buf + (1-β)·grad                  # EMA (identical for both)
+    pre  =  (1-β)·grad + β·buf                   # Nesterov blend (identical for both)
+
+    Muon branch :  update = NS(pre) · scale      # orthogonalise
+    Sign branch :  update = sign(pre) / √n       # sign-normalised
+
+Because the momentum path is fully shared, τ can be freely scheduled — or even
+changed every step — without resetting or corrupting any state.
+
+Frobenius-norm alignment
+------------------------
+Muon's NS output has F-norm ≈ √m for an m×n matrix (m = rows).
+sign(pre) has F-norm = √(m·n).
+Dividing by √n makes the Sign branch produce the same F-norm as Muon,
+so the learning rate carries the same meaning in both branches.
+
+Usage (single device)
+---------------------
+    from stochastic_sign_muon import SingleDeviceStochasticSignMuon
+
+    hidden_weights = [p for p in model.body.parameters() if p.ndim >= 2]
+    rest = (
+        [p for p in model.body.parameters() if p.ndim < 2]
+        + list(model.head.parameters())
+        + list(model.embed.parameters())
+    )
+
+    optimizer = SingleDeviceStochasticSignMuon(
+        [
+            dict(params=hidden_weights, use_muon=True,  lr=0.02,  weight_decay=0.01),
+            dict(params=rest,           use_muon=False, lr=3e-4,  betas=(0.9, 0.95),
+                 weight_decay=0.01),
+        ],
+        tau=0.0,        # start as pure Muon; schedule upward to inject Sign noise
+    )
+
+    # Inside the training loop you can change tau at any time:
+    optimizer.set_tau(0.3)   # 30 % of steps will now be Sign steps
+"""
+
+import torch
+import torch.distributed as dist
+try:
+    from muon import zeropower_via_newtonschulz5, adam_update
+except ImportError:
+    flag("Unable to import muon, consider downloading it via pip install git+https://github.com/KellerJordan/Muon.git to use Muon backed optimizers")
+
+
+# ---------------------------------------------------------------------------
+# Shared update primitives
+# ---------------------------------------------------------------------------
+
+def _nesterov_momentum(
+        grad: torch.Tensor,
+        buf: torch.Tensor,
+        beta: float,
+) -> torch.Tensor:
+    """
+    EMA update followed by a Nesterov combination.
+
+    Modifies `buf` and `grad` in-place (safe inside @torch.no_grad).
+    Returns the Nesterov pre-update vector.
+    """
+    buf.lerp_(grad, 1.0 - beta)  # buf  ←  β·buf + (1-β)·grad
+    return grad.lerp_(buf, beta)  # returns (1-β)·grad + β·buf  [Nesterov]
+
+
+def _muon_branch(pre: torch.Tensor, ns_steps: int = 5) -> torch.Tensor:
+    """
+    Newton-Schulz orthogonalisation + anisotropic scaling.
+    Output F-norm ≈ √m for an m×n pre-update.
+    Returns a tensor of the same shape as `pre`.
+    """
+    u = pre.view(len(pre), -1) if pre.ndim == 4 else pre
+    u = zeropower_via_newtonschulz5(u, steps=ns_steps)
+    u = u * max(1.0, u.size(-2) / u.size(-1)) ** 0.5
+    return u
+
+
+def _sign_branch(pre: torch.Tensor, rand_zero: bool = True) -> torch.Tensor:
+    """
+    Signed update normalised to match Muon's output F-norm.
+
+    sign(pre) has F-norm √(m·n). Dividing by √n yields √m, matching Muon.
+    Zero entries are randomised to ±1 when rand_zero=True.
+    Returns a tensor of the same shape as `pre`.
+    """
+    u = pre.view(len(pre), -1) if pre.ndim == 4 else pre
+
+    s = torch.sign(u)
+
+    if rand_zero:
+        zero_mask = s == 0
+        if zero_mask.any():
+            s[zero_mask] = (
+                    torch.randint(
+                        0, 2,
+                        (zero_mask.sum().item(),),
+                        dtype=u.dtype,
+                        device=u.device,
+                    ) * 2 - 1
+            )
+
+    # Normalise: √(m·n) / √n = √m  →  same F-norm as the Muon branch
+    s = s / (u.size(-1) ** 0.5)
+    return s
+
+
+# ---------------------------------------------------------------------------
+# Single-device variant
+# ---------------------------------------------------------------------------
+
+class SingleDeviceStochasticSignMuon(torch.optim.Optimizer):
+    """
+    Non-distributed stochastic Muon ↔ SignSGD interpolation with AdamW aux.
+
+    Parameters
+    ----------
+    param_groups : list[dict]
+        Each group must contain ``use_muon`` (bool).
+
+        use_muon=True  groups recognise:  lr, momentum, weight_decay
+        use_muon=False groups recognise:  lr, betas, eps, weight_decay
+
+    tau : float
+        Probability of taking a Sign step instead of a Muon step. [0, 1].
+        0.0 → pure Muon (default).  1.0 → pure SignSGD.
+        Schedule via :meth:`set_tau`. No state reset required.
+
+    sign_lr_scale : float
+        Multiplicative downscale applied to lr when the Sign branch fires.
+        Compensates for Sign's lower signal-to-noise ratio vs Muon at equal
+        F-norm. Default 0.1; tune in [0.05, 0.2].
+        Schedule via :meth:`set_sign_lr_scale`. No state reset required.
+        Weight decay is always applied at the base group lr regardless.
+
+    rand_zero : bool
+        Randomise zero-gradient sign entries in the Sign branch (default True).
+
+    nesterov : bool
+        Use Nesterov momentum (default True, matching canonical Muon).
+
+    ns_steps : int
+        Newton-Schulz iteration count for the Muon branch (default 5).
+    """
+
+    def __init__(
+            self,
+            param_groups,
+            tau: float = 0.0,
+            sign_lr_scale: float = 0.1,
+            rand_zero: bool = True,
+            nesterov: bool = True,
+            ns_steps: int = 5,
+    ):
+        assert 0.0 <= tau <= 1.0, f"tau must be in [0, 1], got {tau}"
+        assert sign_lr_scale > 0.0, f"sign_lr_scale must be > 0, got {sign_lr_scale}"
+
+        self.tau = tau
+        self.sign_lr_scale = sign_lr_scale
+        self.rand_zero = rand_zero
+        self.nesterov = nesterov
+        self.ns_steps = ns_steps
+
+        for group in param_groups:
+            assert "use_muon" in group, "Every param group must have a 'use_muon' key."
+            if group["use_muon"]:
+                group.setdefault("lr", 0.02)
+                group.setdefault("momentum", 0.95)
+                group.setdefault("weight_decay", 0.0)
+                allowed = {"params", "lr", "momentum", "weight_decay", "use_muon"}
+                assert set(group.keys()) == allowed, (
+                    f"Unexpected keys in Muon group: {set(group.keys()) - allowed}"
+                )
+            else:
+                group.setdefault("lr", 3e-4)
+                group.setdefault("betas", (0.9, 0.95))
+                group.setdefault("eps", 1e-10)
+                group.setdefault("weight_decay", 0.0)
+                allowed = {"params", "lr", "betas", "eps", "weight_decay", "use_muon"}
+                assert set(group.keys()) == allowed, (
+                    f"Unexpected keys in Adam group: {set(group.keys()) - allowed}"
+                )
+
+        super().__init__(param_groups, {})
+
+    # ------------------------------------------------------------------
+    # Scheduling API
+    # ------------------------------------------------------------------
+
+    def set_tau(self, tau: float) -> None:
+        """
+        Set the Muon ↔ Sign mixing probability.
+        Safe to call at any point; does not reset any optimizer state.
+        """
+        assert 0.0 <= tau <= 1.0, f"tau must be in [0, 1], got {tau}"
+        self.tau = tau
+
+    def set_sign_lr_scale(self, scale: float) -> None:
+        """
+        Set the lr multiplier for Sign steps.
+        Safe to call at any point; does not reset any optimizer state.
+        """
+        assert scale > 0.0, f"sign_lr_scale must be > 0, got {scale}"
+        self.sign_lr_scale = scale
+
+    # ------------------------------------------------------------------
+    # Step
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        # One global Bernoulli draw: all 2-D params follow the same branch this step.
+        use_sign: bool = (torch.rand(1).item() < self.tau)
+
+        for group in self.param_groups:
+
+            if group["use_muon"]:
+                beta = group["momentum"]
+                base_lr = group["lr"]
+                effective_lr = base_lr * self.sign_lr_scale if use_sign else base_lr
+
+                for p in group["params"]:
+                    if p.grad is None:
+                        p.grad = torch.zeros_like(p)
+
+                    state = self.state[p]
+                    if not state:
+                        state["momentum_buffer"] = torch.zeros_like(p)
+
+                    # Shared momentum step (identical regardless of branch)
+                    pre = _nesterov_momentum(p.grad, state["momentum_buffer"], beta)
+
+                    # Branch selection
+                    if use_sign:
+                        update = _sign_branch(pre, rand_zero=self.rand_zero)
+                    else:
+                        update = _muon_branch(pre, ns_steps=self.ns_steps)
+
+                    # Weight decay at base lr; parameter update at effective lr
+                    p.mul_(1.0 - base_lr * group["weight_decay"])
+                    p.add_(update.reshape(p.shape), alpha=-effective_lr)
+
+            else:  # AdamW
+                for p in group["params"]:
+                    if p.grad is None:
+                        p.grad = torch.zeros_like(p)
+
+                    state = self.state[p]
+                    if not state:
+                        state["exp_avg"] = torch.zeros_like(p)
+                        state["exp_avg_sq"] = torch.zeros_like(p)
+                        state["step"] = 0
+
+                    state["step"] += 1
+                    update = adam_update(
+                        p.grad,
+                        state["exp_avg"],
+                        state["exp_avg_sq"],
+                        state["step"],
+                        group["betas"],
+                        group["eps"],
+                    )
+                    p.mul_(1.0 - group["lr"] * group["weight_decay"])
+                    p.add_(update, alpha=-group["lr"])
+
+        return loss
+
+
+# ---------------------------------------------------------------------------
+# Distributed variant
+# ---------------------------------------------------------------------------
+
+class StochasticSignMuon(torch.optim.Optimizer):
+    """
+    Distributed stochastic Muon ↔ SignSGD interpolation with AdamW aux.
+
+    Requires an initialised ``torch.distributed`` process group.
+    Parameters are sorted by size and sharded across ranks exactly as in the
+    original ``MuonWithAuxAdam``.
+
+    The Bernoulli draw is made on rank 0 and broadcast to all ranks so that
+    every GPU takes the same branch on every step.
+
+    See ``SingleDeviceStochasticSignMuon`` for the full parameter docstring.
+    """
+
+    def __init__(
+            self,
+            param_groups,
+            tau: float = 0.0,
+            sign_lr_scale: float = 0.1,
+            rand_zero: bool = True,
+            nesterov: bool = True,
+            ns_steps: int = 5,
+    ):
+        assert 0.0 <= tau <= 1.0, f"tau must be in [0, 1], got {tau}"
+        assert sign_lr_scale > 0.0, f"sign_lr_scale must be > 0, got {sign_lr_scale}"
+
+        self.tau = tau
+        self.sign_lr_scale = sign_lr_scale
+        self.rand_zero = rand_zero
+        self.nesterov = nesterov
+        self.ns_steps = ns_steps
+
+        for group in param_groups:
+            assert "use_muon" in group, "Every param group must have a 'use_muon' key."
+            if group["use_muon"]:
+                group["params"] = sorted(
+                    group["params"], key=lambda x: x.size(), reverse=True
+                )
+                group.setdefault("lr", 0.02)
+                group.setdefault("momentum", 0.95)
+                group.setdefault("weight_decay", 0.0)
+                allowed = {"params", "lr", "momentum", "weight_decay", "use_muon"}
+                assert set(group.keys()) == allowed
+            else:
+                group.setdefault("lr", 3e-4)
+                group.setdefault("betas", (0.9, 0.95))
+                group.setdefault("eps", 1e-10)
+                group.setdefault("weight_decay", 0.0)
+                allowed = {"params", "lr", "betas", "eps", "weight_decay", "use_muon"}
+                assert set(group.keys()) == allowed
+
+        super().__init__(param_groups, {})
+
+    # ------------------------------------------------------------------
+    # Scheduling API
+    # ------------------------------------------------------------------
+
+    def set_tau(self, tau: float) -> None:
+        """Schedule the Muon ↔ Sign mixing probability. Safe from any rank."""
+        assert 0.0 <= tau <= 1.0, f"tau must be in [0, 1], got {tau}"
+        self.tau = tau
+
+    def set_sign_lr_scale(self, scale: float) -> None:
+        """Schedule the lr multiplier for Sign steps. Safe from any rank."""
+        assert scale > 0.0, f"sign_lr_scale must be > 0, got {scale}"
+        self.sign_lr_scale = scale
+
+    # ------------------------------------------------------------------
+    # Step
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+
+        # Broadcast the Bernoulli draw from rank 0 so every GPU takes the same branch.
+        use_sign_t = torch.zeros(1, dtype=torch.float32,
+                                 device=torch.device("cuda", rank))
+        if rank == 0:
+            use_sign_t[0] = 1.0 if torch.rand(1).item() < self.tau else 0.0
+        dist.broadcast(use_sign_t, src=0)
+        use_sign: bool = use_sign_t.item() > 0.5
+
+        for group in self.param_groups:
+
+            if group["use_muon"]:
+                beta = group["momentum"]
+                base_lr = group["lr"]
+                effective_lr = base_lr * self.sign_lr_scale if use_sign else base_lr
+
+                params = group["params"]
+                params_pad = params + [torch.empty_like(params[-1])] * (
+                        world_size - len(params) % world_size
+                )
+
+                for base_i in range(0, len(params), world_size):
+                    local_idx = base_i + rank
+                    if local_idx < len(params):
+                        p = params[local_idx]
+                        if p.grad is None:
+                            p.grad = torch.zeros_like(p)
+
+                        state = self.state[p]
+                        if not state:
+                            state["momentum_buffer"] = torch.zeros_like(p)
+
+                        pre = _nesterov_momentum(
+                            p.grad, state["momentum_buffer"], beta
+                        )
+
+                        if use_sign:
+                            update = _sign_branch(pre, rand_zero=self.rand_zero)
+                        else:
+                            update = _muon_branch(pre, ns_steps=self.ns_steps)
+
+                        # Weight decay at base lr; parameter update at effective lr
+                        p.mul_(1.0 - base_lr * group["weight_decay"])
+                        p.add_(update.reshape(p.shape), alpha=-effective_lr)
+
+                    dist.all_gather(
+                        params_pad[base_i: base_i + world_size],
+                        params_pad[base_i + rank],
+                    )
+
+            else:  # AdamW
+                for p in group["params"]:
+                    if p.grad is None:
+                        p.grad = torch.zeros_like(p)
+
+                    state = self.state[p]
+                    if not state:
+                        state["exp_avg"] = torch.zeros_like(p)
+                        state["exp_avg_sq"] = torch.zeros_like(p)
+                        state["step"] = 0
+
+                    state["step"] += 1
+                    update = adam_update(
+                        p.grad,
+                        state["exp_avg"],
+                        state["exp_avg_sq"],
+                        state["step"],
+                        group["betas"],
+                        group["eps"],
+                    )
+                    p.mul_(1.0 - group["lr"] * group["weight_decay"])
+                    p.add_(update, alpha=-group["lr"])
+
+        return loss

vathos/training.py

@@ -0,0 +1,129 @@
+from tqdm.notebook import tqdm
+from Vathos.blocks import *
+
+
+def symbolic_1d_ar_target_train(model, train_loader, val_loader, optimizer, scheduler, criterion,
+          device, NUM_EPOCHS=100, use_amp=False, CLIP_NORM=1.0, div=1):
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+
+    steps_per_epoch = len(train_loader) // div
+
+    for epoch in range(NUM_EPOCHS):
+        model.train()
+        running_loss = 0.0
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{NUM_EPOCHS} Train", total=steps_per_epoch)
+        steps = 0
+
+        for batch_x, batch_tgt in pbar:
+            batch_x = batch_x.to(device)
+            steps += 1
+
+            optimizer.zero_grad()
+            with torch.amp.autocast('cuda'):
+                logits = model(batch_x)  # [B,S,V]
+                loss = criterion(
+                    logits[:, :-1, :].reshape(-1, logits.size(-1)),
+                    batch_x[:, 1:].reshape(-1)
+                )
+
+            scaler.scale(loss).backward()
+
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP_NORM)
+            scaler.step(optimizer)
+            scaler.update()
+
+            scheduler.step()
+
+            running_loss += loss.item()
+            lr_now = optimizer.param_groups[0]['lr']
+            pbar.set_postfix({'loss': running_loss / (pbar.n + 1), 'lr': lr_now})
+
+            if steps >= steps_per_epoch:
+                break
+
+        model.eval()
+        val_loss = 0.0
+        # with torch.no_grad():
+        if True:
+            for batch_x, batch_tgt, composer_idx in tqdm(val_loader, desc=f"Epoch {epoch + 1}/{NUM_EPOCHS} Val"):
+                batch_x = batch_x.to(device).requires_grad_(False)
+
+                logits = model(batch_x)
+                loss = criterion(
+                    logits[:, :-1, :].reshape(-1, logits.size(-1)),
+                    batch_x[:, 1:].reshape(-1)
+                )
+                val_loss += loss.item()
+                if isinstance(model, VathosModel):
+                    model.module.register_loss(loss.item())
+
+        val_avg = val_loss / len(val_loader)
+        if isinstance(model, VathosModel):
+            model.module.register_epoch()
+        print(f"Epoch {epoch + 1} val avg loss: {val_avg:.4f}")
+
+
+def symbolic_1d_ar_input_ids_train(model, train_loader, val_loader, optimizer, scheduler, criterion,
+                                      device, NUM_EPOCHS=100, use_amp=False, CLIP_NORM=1.0, spe=1000, val_steps=float('inf')):
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+
+    steps_per_epoch = spe
+
+    for epoch in range(NUM_EPOCHS):
+        model.train()
+        running_loss = 0.0
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{NUM_EPOCHS} Train", total=steps_per_epoch)
+        steps = 0
+
+        for batch_x in pbar:
+            batch_x = batch_x['input_ids'].to(device)
+            steps += 1
+
+            optimizer.zero_grad()
+            with torch.amp.autocast('cuda'):
+                logits = model(batch_x)  # [B,S,V]
+                loss = criterion(
+                    logits[:, :-1, :].reshape(-1, logits.size(-1)),
+                    batch_x[:, 1:].reshape(-1)
+                )
+
+            scaler.scale(loss).backward()
+
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP_NORM)
+            scaler.step(optimizer)
+            scaler.update()
+            if scheduler is not None:
+                scheduler.step()
+
+            running_loss += loss.item()
+            lr_now = optimizer.param_groups[0]['lr']
+            pbar.set_postfix({'loss': running_loss / (pbar.n + 1), 'lr': lr_now})
+
+            if steps >= steps_per_epoch:
+                break
+
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+        # if True:
+            val_s = 0
+            for batch_x in tqdm(val_loader, desc=f"Epoch {epoch + 1}/{NUM_EPOCHS} Val"):
+                batch_x = batch_x['input_ids'].to(device).requires_grad_(False)
+                val_s += 1
+                if val_s >= val_steps:
+                    break
+                logits = model(batch_x)
+                loss = criterion(
+                    logits[:, :-1, :].reshape(-1, logits.size(-1)),
+                    batch_x[:, 1:].reshape(-1)
+                )
+                model.module.register_loss(loss.item())
+
+        val_avg = model.module.get_mean_loss()
+        model.module.register_epoch()
+        print(f"Epoch {epoch + 1} val avg loss: {val_avg:.4f}")
+
+
+