Publish paper2code skill to marketplace

Constraint: Public skills are published only by explicit administrator action unless they are tracked third-party market sources.
Confidence: high
Scope-risk: narrow
Directive: Keep private/internal skills out of the public marketplace.
Tested: Marketplace validation passed.

.agents/plugins/marketplace.json

@@ -111,6 +111,18 @@
         "authentication": "ON_INSTALL"
       },
       "category": "文档处理"
+    },
+    {
+      "name": "paper2code",
+      "source": {
+        "source": "local",
+        "path": "./plugins/codex/plugins/paper2code"
+      },
+      "policy": {
+        "installation": "AVAILABLE",
+        "authentication": "ON_INSTALL"
+      },
+      "category": "文档处理"
     }
   ]
 }

plugins/codex/.agents/plugins/marketplace.json

@@ -111,6 +111,18 @@
         "authentication": "ON_INSTALL"
       },
       "category": "文档处理"
+    },
+    {
+      "name": "paper2code",
+      "source": {
+        "source": "local",
+        "path": "./plugins/paper2code"
+      },
+      "policy": {
+        "installation": "AVAILABLE",
+        "authentication": "ON_INSTALL"
+      },
+      "category": "文档处理"
     }
   ]
 }

plugins/codex/plugins/paper2code/.codex-plugin/plugin.json

@@ -0,0 +1,37 @@
+{
+  "name": "paper2code",
+  "version": "0.1.0",
+  "description": "将Arxiv论文转换为简洁的、基于引用的Python实现代码。当用户使用`/paper2code`命令并提供Arxiv链接或论文ID时，该系统会提示“实现这篇论文”；或者当用户粘贴Arxiv链接请求代码实现时，系统也会作出相应响应。系统会如实指出所有可能存在歧义的地方，绝不会添加论文中未提及的实现细节。",
+  "author": {
+    "name": "EAPIL",
+    "url": "https://git.playones.com/arechen/EapilSkillMarket"
+  },
+  "homepage": "https://git.playones.com/arechen/EapilSkillMarket",
+  "repository": "https://git.playones.com/arechen/EapilSkillMarket",
+  "license": "Proprietary",
+  "keywords": [
+    "eapil",
+    "codex-skill",
+    "paper2code"
+  ],
+  "skills": "./skills/",
+  "interface": {
+    "displayName": "paper2code",
+    "shortDescription": "将Arxiv论文转换为简洁的、基于引用的Python实现代码。当用户使用`/paper2code`命令并提供Arxiv链接或论文ID时，该系统会提示“实现这篇论文”；或者当用户粘贴Arxiv链接请求代码实现时，系统也会作出相应响应。系统会如实指出所有可能存在歧义的地方，绝不会添加论文中未提及的实现细节。",
+    "longDescription": "将Arxiv论文转换为简洁的、基于引用的Python实现代码。当用户使用`/paper2code`命令并提供Arxiv链接或论文ID时，该系统会提示“实现这篇论文”；或者当用户粘贴Arxiv链接请求代码实现时，系统也会作出相应响应。系统会如实指出所有可能存在歧义的地方，绝不会添加论文中未提及的实现细节。",
+    "developerName": "EAPIL",
+    "category": "文档处理",
+    "capabilities": [
+      "Read",
+      "Write"
+    ],
+    "defaultPrompt": [
+      "使用 paper2code 帮我完成这个任务。"
+    ],
+    "websiteURL": "https://git.playones.com/arechen/EapilSkillMarket",
+    "privacyPolicyURL": "https://git.playones.com/arechen/EapilSkillMarket",
+    "termsOfServiceURL": "https://git.playones.com/arechen/EapilSkillMarket",
+    "brandColor": "#2563EB",
+    "screenshots": []
+  }
+}

plugins/codex/plugins/paper2code/skills/paper2code/SKILL.md

@@ -0,0 +1,112 @@
+---
+name: paper2code
+description: Converts an arxiv paper into a minimal, citation-anchored Python implementation. Trigger when user runs /paper2code with an arxiv URL or paper ID, says "implement this paper", or pastes an arxiv link asking for implementation. Flags all ambiguities honestly. Never invents implementation details not stated in the paper.
+---
+
+# paper2code — Orchestration
+
+You are executing the paper2code skill. This file governs the high-level flow. Each stage dispatches to a detailed reasoning protocol in `pipeline/`. Do NOT skip stages. Do NOT combine stages. Execute them in order.
+
+## Parse arguments
+
+Extract from the user's input:
+- `ARXIV_ID`: the arxiv paper ID (e.g., `2106.09685`). Strip any URL prefix.
+- `MODE`: one of `minimal` (default), `full`, `educational`.
+- `FRAMEWORK`: one of `pytorch` (default), `jax`, `numpy`.
+
+If the user provided a full URL like `https://arxiv.org/abs/2106.09685`, extract the ID `2106.09685`.
+If the user provided a versioned ID like `2106.09685v2`, keep the version.
+
+## Set up working directory
+
+Create a temporary working directory: `.paper2code_work/{ARXIV_ID}/`
+This is where intermediate artifacts go. The final output goes in the current directory under `{paper_slug}/`.
+
+## Install dependencies
+
+Run via Bash:
+```bash
+pip install pymupdf4llm pdfplumber requests pyyaml
+```
+
+## Execute pipeline
+
+### Stage 1 — Paper Acquisition and Parsing
+Read and follow: `pipeline/01_paper_acquisition.md`
+
+Run the helper script to fetch and parse the paper:
+```bash
+python skills/paper2code/scripts/fetch_paper.py {ARXIV_ID} .paper2code_work/{ARXIV_ID}/
+```
+Then run structure extraction:
+```bash
+python skills/paper2code/scripts/extract_structure.py .paper2code_work/{ARXIV_ID}/paper_text.md .paper2code_work/{ARXIV_ID}/
+```
+Verify the outputs exist before proceeding. If extraction failed, follow the fallback protocol in `pipeline/01_paper_acquisition.md`.
+
+The script also searches for official code repositories (in the paper text and on the arxiv page) and saves any found links to `paper_metadata.json` under the `official_code` key. Verify these links before relying on them — see Step 8 in `pipeline/01_paper_acquisition.md`.
+
+### Stage 2 — Contribution Identification
+Read and follow: `pipeline/02_contribution_identification.md`
+
+Read the parsed paper sections. Identify the single core contribution. Classify the paper type. Write the contribution statement. Save it to `.paper2code_work/{ARXIV_ID}/contribution.md`.
+
+### Stage 3 — Ambiguity Audit
+Read and follow: `pipeline/03_ambiguity_audit.md`
+
+Before reading this stage, also read: `guardrails/hallucination_prevention.md`
+
+Go through every implementation-relevant detail. Classify each as SPECIFIED, PARTIALLY_SPECIFIED, or UNSPECIFIED. Save the audit to `.paper2code_work/{ARXIV_ID}/ambiguity_audit.md`.
+
+### Stage 4 — Code Generation
+Read and follow: `pipeline/04_code_generation.md`
+
+Before writing code, read:
+- `guardrails/scope_enforcement.md` — to determine what's in and out of scope
+- `guardrails/badly_written_papers.md` — if the paper is vague or inconsistent
+- The relevant knowledge files in `knowledge/` for the paper's domain
+- The scaffold templates in `scaffolds/` for the expected file structure
+
+Determine the `paper_slug` from the paper title (lowercase, underscores, no special chars).
+Generate all files under `{paper_slug}/` in the current working directory.
+
+### Stage 5 — Walkthrough Notebook
+Read and follow: `pipeline/05_walkthrough_notebook.md`
+
+Generate the walkthrough notebook that connects paper sections to code with runnable sanity checks. Save to `{paper_slug}/notebooks/walkthrough.ipynb`.
+
+## Cleanup
+
+Remove the `.paper2code_work/` directory after successful completion.
+
+## Final output
+
+Print a summary:
+```
+✓ paper2code complete for: {paper_title}
+  Output directory: {paper_slug}/
+  Files generated: {list of files}
+  Unspecified choices: {count} (see REPRODUCTION_NOTES.md)
+  Mode: {MODE} | Framework: {FRAMEWORK}
+```
+
+## Mode-specific behavior
+
+- **minimal** (default): Core contribution only. Training loop only if contribution involves training. No data pipeline beyond Dataset skeleton.
+- **full**: Core contribution + full training loop + data pipeline + evaluation pipeline. More code, same citation rigor.
+- **educational**: Same as minimal but with extra inline comments explaining ML concepts, expanded walkthrough notebook with theory sections, and a `PAPER_GUIDE.md` that walks through the paper section by section.
+
+## Guardrails — always active
+
+These apply at ALL stages. Read them if you haven't already:
+- `guardrails/hallucination_prevention.md` — the most important file in this skill
+- `guardrails/scope_enforcement.md` — what to implement and what to skip
+- `guardrails/badly_written_papers.md` — what to do when the paper is unclear
+
+## Knowledge base — consult as needed
+
+Before implementing any of these components, read the corresponding knowledge file:
+- Transformer layers, attention, positional encoding → `knowledge/transformer_components.md`
+- Optimizers, LR schedules, batch size semantics → `knowledge/training_recipes.md`
+- Cross-entropy, contrastive loss, diffusion loss, ELBO → `knowledge/loss_functions.md`
+- Framework-specific pitfalls, notation mismatches → `knowledge/paper_to_code_mistakes.md`

plugins/codex/plugins/paper2code/skills/paper2code/guardrails/badly_written_papers.md

@@ -0,0 +1,215 @@
+# Guardrail: Handling Badly Written Papers
+
+## Reality check
+
+Many papers are vague, inconsistent, or incomplete. This is not a criticism — writing a paper is hard, page limits are strict, and reviewers rarely check hyperparameter tables. But it means you will regularly encounter papers where the text alone is insufficient to produce a correct implementation.
+
+This file tells you what to do when that happens. The answer is never "guess silently."
+
+---
+
+## Decision tree for resolving ambiguity
+
+```
+Is the detail stated explicitly in the paper?
+├── YES → Use it. Cite the section. Done.
+├── PARTIALLY → Follow the partial specification protocol (below)
+└── NO →
+    Is there official code on GitHub?
+    ├── YES → Use it as ground truth. Cite the file and line number.
+    │         Mark as [FROM_OFFICIAL_CODE], not [SPECIFIED].
+    └── NO →
+        Is there a well-known reimplementation?
+        ├── YES → Reference it in REPRODUCTION_NOTES.md.
+        │         DO NOT blindly copy — they made their own choices.
+        │         Mark as [UNSPECIFIED] and note the external reference.
+        └── NO →
+            Is there a "standard" choice in the field?
+            ├── YES → Use it. Mark as [UNSPECIFIED] with alternatives.
+            └── NO →
+                Write a stub with a detailed docstring.
+                Explain what should go here and what the options are.
+```
+
+---
+
+## Partial specification protocol
+
+When the paper partially specifies something:
+
+### "We use Adam"
+The paper says an optimizer but not all parameters.
+```python
+# §4.1 — "We use the Adam optimizer"
+# [PARTIALLY_SPECIFIED] Optimizer stated as Adam, but beta and epsilon values not specified
+# Using: β₁=0.9, β₂=0.999, ε=1e-8 (PyTorch defaults)
+# Alternatives: β₁=0.9, β₂=0.98, ε=1e-9 (Transformer default from Vaswani et al.)
+optimizer = torch.optim.Adam(model.parameters(), lr=config.lr,
+                              betas=(0.9, 0.999), eps=1e-8)
+```
+
+### "We use dropout for regularization"
+The paper mentions dropout but not the rate or placement.
+```python
+# §3.3 — "We use dropout for regularization"
+# [PARTIALLY_SPECIFIED] Dropout mentioned but rate not specified
+# Using: 0.1 (common default for transformer models)
+# Placement: after attention and FFN (not specified — this is our choice)
+self.dropout = nn.Dropout(0.1)
+```
+
+### "Similar to Transformer [Vaswani et al., 2017]"
+The paper references another work for a component.
+```python
+# §3.1 — "We use an architecture similar to the Transformer [Vaswani et al., 2017]"
+# [PARTIALLY_SPECIFIED] "Similar" implies differences exist but none are specified
+# We implement the standard Transformer from Vaswani et al. unless other sections
+# describe modifications. Reader should check if the authors intended differences.
+```
+
+---
+
+## When the paper contradicts itself
+
+### Figure vs text
+Figures are often created early in the paper-writing process and may show an earlier version of the architecture or method. Text is usually more up-to-date.
+
+**Rule: Trust the text. Flag the figure discrepancy.**
+```python
+# §3.2 — "We apply layer normalization before each sub-layer" (Pre-LN)
+# NOTE: Figure 2 shows post-norm placement, inconsistent with this text.
+# We implement pre-norm as stated in the text. If reproduction fails,
+# try post-norm (change this to: x = self.norm(x + sublayer(x)))
+```
+
+### Equation vs prose
+Equations are peer-reviewed more carefully and are more precise.
+
+**Rule: Trust the equation. Flag the prose discrepancy.**
+
+### Different numbers in different sections
+The paper might say "learning rate of 1e-4" in one section and "learning rate of 3e-4" in another.
+
+**Rule: If one is in the appendix/hyperparameter table, trust that. If both are in prose, flag both and use the one from the main experimental section (usually the one paired with the best results).**
+
+---
+
+## When the paper is genuinely incomplete
+
+Some papers are missing critical details that make reproduction impossible without additional information. Here's what to do:
+
+### Missing architecture details (for architecture papers)
+If you cannot determine the model architecture from the paper:
+1. Check for official code — this is the only reliable resolution
+2. Check for well-known reimplementations (HuggingFace, lucidrains, etc.)
+3. If neither exists, implement what you CAN determine and write stubs for the rest:
+
+```python
+def mysterious_component(x: torch.Tensor) -> torch.Tensor:
+    """§3.4 — The paper describes a 'novel aggregation mechanism' but does not
+    provide sufficient detail to implement it.
+    
+    What we know:
+      - Input: tensor of shape (batch, seq_len, d_model) — from §3.3
+      - Output: tensor of shape (batch, d_model) — inferred from §3.5 which uses the output
+      - It "aggregates information across the sequence" — from §3.4
+    
+    What we don't know:
+      - The specific aggregation operation (mean pooling? attention? learned query?)
+      - Whether it has learnable parameters
+      - Whether it uses masking
+    
+    To complete this implementation, check:
+      1. The authors' official code (if released after this generation)
+      2. Follow-up papers that reproduce this work
+      3. Contact the authors
+    """
+    raise NotImplementedError(
+        "Cannot implement: paper §3.4 does not provide sufficient architectural detail. "
+        "See docstring for what is known and what is missing."
+    )
+```
+
+### Missing training details (for training method papers)
+If the training procedure is the contribution but is underspecified:
+1. Implement the parts that ARE specified
+2. For unspecified parts, use the most common baseline and flag it loudly
+3. Add a section in REPRODUCTION_NOTES.md titled "Critical unspecified training details"
+
+### Missing evaluation details
+Note which metric version you're using and that it might not match:
+```python
+# [UNSPECIFIED] Paper reports "BLEU score" but does not specify which implementation
+# Using: sacrebleu (the de facto standard for reproducible BLEU)
+# NOTE: Different BLEU implementations can give significantly different numbers
+# (sometimes 1-2 points) due to tokenization and smoothing differences
+```
+
+---
+
+## Using official code as ground truth
+
+When official code exists, use it to resolve ambiguities. But be careful:
+
+### DO:
+- Link to the exact file and line number
+- Mark resolved items as `[FROM_OFFICIAL_CODE]`
+- Note when the official code differs from the paper text
+- Check the commit history — early commits are closer to the paper, later commits may reflect improvements
+
+### DON'T:
+- Assume official code is always right — it may have bugs
+- Copy code style or non-essential patterns
+- Use code from Pull Requests or branches — stick to main/master
+- Assume the code matches the paper — sometimes best results come from later improvements not in the paper
+
+### Format:
+```python
+# [FROM_OFFICIAL_CODE] Weight initialization: normal with std=0.02
+# Source: github.com/author/repo/blob/main/model.py#L42
+# Paper does not specify initialization (§3.1)
+nn.init.normal_(self.weight, std=0.02)
+```
+
+---
+
+## Using well-known reimplementations
+
+HuggingFace, lucidrains, Andrej Karpathy, Phil Wang — some implementers are widely trusted. When using their implementations as reference:
+
+1. **Reference in REPRODUCTION_NOTES.md**, not in inline comments (unless resolving a specific choice)
+2. **Note that they may have made their own unspecified choices** — their choices are educated guesses, not specifications
+3. **Do not assume they are correct** — well-known reimplementations sometimes have bugs or intentional deviations
+4. **Use them for sanity checking**, not as ground truth
+
+---
+
+## Known erratum handling
+
+If the paper has a published erratum, correction, or the authors have acknowledged an error:
+
+```python
+# §3.2, Eq. 4 — Original paper has a typo: denominator should be sqrt(d_k), not d_k
+# ERRATUM: Acknowledged by authors in arxiv v2 (see changelog)
+# We implement the corrected version
+scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+```
+
+Check arxiv for paper versions — v2, v3 often contain corrections. The abstract page shows the version history.
+
+---
+
+## When to give up
+
+There is a point where a paper is so underspecified that a meaningful implementation is impossible. Signs:
+
+- The core method relies on a component described only as "proprietary" or "internal"
+- The paper describes results but not the method (sometimes in position papers)
+- The paper is a survey or review with no novel method
+- The paper's contribution is purely theoretical with no implementable component
+
+In these cases:
+1. Generate what you CAN (data loading, evaluation metrics, utility functions)
+2. Write a clear `REPRODUCTION_NOTES.md` explaining why a full implementation is not possible
+3. List exactly what information would be needed to complete it
+4. Do NOT generate fake implementations to make the output look complete

plugins/codex/plugins/paper2code/skills/paper2code/guardrails/hallucination_prevention.md

@@ -0,0 +1,194 @@
+# Guardrail: Hallucination Prevention
+
+## The core problem
+
+The single most dangerous failure mode of paper2code is **confident invention** — generating code that looks plausible but implements details the paper never specified, without flagging them. A researcher who trusts this output will waste days debugging differences that exist because the model guessed.
+
+This file exists to prevent that. Read it before every code generation stage. Internalize it.
+
+---
+
+## The bright line rule
+
+**If a detail is not explicitly stated in the paper text, the appendix, the paper's official GitHub repository, or a well-known replication paper — it is UNSPECIFIED.**
+
+Not "probably this." Not "standard practice." Not "everyone uses." **UNSPECIFIED.**
+
+---
+
+## What counts as "stated in the paper"
+
+### Counts as SPECIFIED:
+- Direct statement: "We use d_model = 512" → SPECIFIED
+- Table entry: Table 3 shows "learning rate: 3e-4" → SPECIFIED
+- Equation: Eq. 4 defines the loss function → SPECIFIED (for the equation, not for implementation details like numerical precision)
+- Algorithm box: Algorithm 1 lines 3-5 describe the update rule → SPECIFIED
+- Footnote: Footnote 3 says "we clip gradients at 1.0" → SPECIFIED
+- Appendix: Appendix B Table 6 shows "warmup steps = 4000" → SPECIFIED
+
+### Does NOT count as SPECIFIED:
+- "We use standard optimization" → **UNSPECIFIED** (standard according to whom? which optimizer? what hyperparameters?)
+- "Following prior work [23]" → **PARTIALLY_SPECIFIED** (you must look up [23] and report what it says, or flag that the reader needs to)
+- "We use Adam" → **PARTIALLY_SPECIFIED** (Adam has β₁, β₂, ε parameters — are they stated?)
+- "Similar architecture to [X]" → **PARTIALLY_SPECIFIED** (similar is not identical — what differs?)
+- "Standard hyperparameters" → **UNSPECIFIED**
+- "Default settings" → **UNSPECIFIED** (whose defaults? PyTorch? TensorFlow? They differ.)
+- Descriptions in related work of other people's methods → **NOT A SPECIFICATION OF THIS PAPER'S METHOD**
+- Blog posts, tweets, or talks by the authors → **NOT PEER-REVIEWED, note as supplementary only**
+
+---
+
+## The UNSPECIFIED comment protocol
+
+When you make a choice for an UNSPECIFIED item, the code comment must have three parts:
+
+```python
+# [UNSPECIFIED] {What the paper doesn't specify}
+# Using: {your choice}
+# Alternatives: {other reasonable choices}
+```
+
+Example:
+```python
+# [UNSPECIFIED] Paper does not state activation function in the feed-forward network
+# Using: GELU (most common in recent transformer implementations)
+# Alternatives: ReLU (original transformer), SiLU/Swish (used in LLaMA, PaLM)
+self.activation = nn.GELU()
+```
+
+You must NEVER write just:
+```python
+self.activation = nn.GELU()  # standard choice
+```
+This hides the fact that the paper didn't specify it.
+
+---
+
+## Equation ground truth rule
+
+Equations are more precise than prose. If the prose description and the equation conflict:
+
+1. Implement the equation
+2. Flag the discrepancy explicitly:
+
+```python
+# §3.2, Eq. 4 — loss = -log(exp(sim(z_i, z_j)/τ) / Σ_k exp(sim(z_i, z_k)/τ))
+# NOTE: The prose in §3.2 says "we average over all positive pairs" but Eq. 4
+# sums rather than averages. We implement Eq. 4 as written.
+```
+
+This applies even when the equation is clearly a typo. Implement the equation, flag the likely typo, and add a comment about what the authors probably meant.
+
+---
+
+## The "standard" trap
+
+Papers frequently use the word "standard" without definition. Here's what to do:
+
+| Paper says | What to do |
+|-----------|------------|
+| "standard transformer" | Ask: which transformer? Pre-norm? Post-norm? How many layers? Flag as UNSPECIFIED unless the paper cites a specific architecture |
+| "standard augmentation" | Ask: which augmentations? Random crop size? Flip probability? Color jitter parameters? Flag as UNSPECIFIED |
+| "standard preprocessing" | Ask: what tokenizer? What normalization? What sequence length? Flag as UNSPECIFIED |
+| "standard evaluation" | Ask: which metric implementation? What post-processing? Flag as UNSPECIFIED |
+| "we follow standard practice" | This means nothing. Flag as UNSPECIFIED. |
+
+---
+
+## The number precision trap
+
+Do not silently change numbers:
+
+- If the paper says 512, use 512 — not 256 "for simplicity" without flagging it
+- If the paper says 0.9, use 0.9 — not 0.99 because "that's what people usually use"
+- If the paper says 100k steps, use 100000 — not 50000 because "it should converge faster"
+
+Any deviation from stated numbers must be flagged:
+```python
+# §5.2 states d_model = 512, but this walkthrough uses d_model = 64 for CPU execution
+# Set d_model = 512 for actual reproduction
+```
+
+---
+
+## The initialization trap
+
+Weight initialization is almost never specified in papers but matters enormously. If the paper does not specify initialization:
+
+```python
+# [UNSPECIFIED] Paper does not describe weight initialization
+# Using: PyTorch defaults (Kaiming uniform for Linear, uniform for Embedding)
+# Alternatives: Xavier uniform (common for transformers), normal init with std=0.02
+# NOTE: Initialization can significantly affect training stability and convergence
+```
+
+If the paper specifies initialization (rare but valuable), implement it exactly and cite the section.
+
+---
+
+## The framework translation trap
+
+Different frameworks have different defaults that papers don't always clarify:
+
+### Batch normalization momentum
+- Paper says `momentum = 0.1`
+- PyTorch BatchNorm uses `momentum = 0.1` (but its definition is inverted: `running_mean = (1 - momentum) * running_mean + momentum * batch_mean`)
+- TensorFlow BatchNorm uses `momentum = 0.99` (with convention: `running_mean = momentum * running_mean + (1 - momentum) * batch_mean`)
+- PyTorch `momentum=0.1` ≈ TensorFlow `momentum=0.9`
+- **Always clarify which convention the paper uses**
+
+### Dropout rate vs keep probability
+- Paper says "dropout 0.1" — does this mean drop probability = 0.1 or keep probability = 0.1?
+- Almost always means drop probability = 0.1 (keep = 0.9), but older papers sometimes use keep probability
+- PyTorch `nn.Dropout(p=0.1)` means drop probability = 0.1
+
+### Layer normalization epsilon
+- Papers almost never specify this
+- PyTorch default: 1e-5
+- Common in papers: 1e-6
+- Some implementations: 1e-8
+- **Flag as UNSPECIFIED and note which you chose**
+
+---
+
+## The "we found" trap
+
+When a paper says "we found that X works better," this is usually an empirical observation, not a derived result. Treat it as useful information but note that:
+- "Better" often means "better for our specific setup/dataset/scale"
+- The alternative (not-X) might work fine for different scenarios
+- This is an [ASSUMPTION] not a [SPECIFIED] item when you implement it as the default
+
+---
+
+## Prohibited phrases in generated code
+
+Never write any of these in code comments:
+- "standard practice" (without specifying what practice)
+- "as usual" (usual for whom?)
+- "obviously" (if it's obvious, you don't need to say it; and it's probably not obvious)
+- "typically" (without a citation)
+- "it's well known that" (without a citation)
+- "for simplicity" (as justification for deviating from the paper)
+- "should work" (either it's specified or it isn't)
+
+---
+
+## The official code shortcut
+
+If official code exists:
+1. Note its URL in REPRODUCTION_NOTES.md
+2. You may use it to resolve UNSPECIFIED items — but:
+   - Mark them as `[FROM_OFFICIAL_CODE]`, not as `SPECIFIED`
+   - The official code may differ from the paper (bug fixes, improvements, errors)
+   - Link to the exact line: `github.com/author/repo/blob/main/model.py#L42`
+3. Do NOT copy-paste code. Read the official code to understand the choice, then implement it yourself with a citation
+
+---
+
+## Self-audit questions
+
+Before finishing any code generation, ask yourself:
+1. If I removed all my `[UNSPECIFIED]` comments, would a reader think the paper specified everything? If yes, I'm probably missing flags.
+2. Did I add any implementation detail from my own ML knowledge without checking if the paper says it? Flag it.
+3. Would the authors of this paper agree that my code matches their description? If I'm not sure, something needs a flag.
+4. Is there a single magic number anywhere without a citation or `[UNSPECIFIED]` comment? Find it and fix it.

plugins/codex/plugins/paper2code/skills/paper2code/guardrails/scope_enforcement.md

@@ -0,0 +1,186 @@
+# Guardrail: Scope Enforcement
+
+## Purpose
+Determine what to implement and what to skip. Implementing too much adds noise. Implementing too little is useless. This decision tree makes scope decisions explicit and traceable.
+
+---
+
+## The decision tree
+
+For each component you're considering implementing, walk through this tree:
+
+```
+Is this component the paper's core contribution?
+├── YES → ALWAYS IN SCOPE. Implement fully with citation anchoring.
+│
+└── NO →
+    Does the paper describe this component in detail (equations, pseudocode, algorithm box)?
+    ├── YES →
+    │   Is this component necessary to understand the core contribution?
+    │   ├── YES → IN SCOPE. Implement with citations.
+    │   └── NO → OUT OF SCOPE. Note its existence, don't implement.
+    │
+    └── NO →
+        Does the paper reference another paper for this component?
+        ├── YES → REFERENCE ONLY. Import or note the dependency.
+        │         Add a comment: "§X.Y — uses {component} from {reference}"
+        │
+        └── NO →
+            Is this a standard ML component (optimizer, standard layer, etc.)?
+            ├── YES → USE FRAMEWORK DEFAULT. Import from PyTorch/library.
+            │         Don't reimplement. Document which version you're using.
+            │
+            └── NO →
+                You've found something the paper assumes you know but doesn't explain.
+                Flag it in REPRODUCTION_NOTES.md and use a reasonable default.
+```
+
+---
+
+## Scope by paper type
+
+### Type (a): New architecture paper
+
+| Component | In scope? | Notes |
+|-----------|-----------|-------|
+| Model architecture | ✅ ALWAYS | This IS the contribution |
+| Forward pass | ✅ ALWAYS | Demonstrates the architecture |
+| Loss function | ⚠️ CONDITIONAL | Only if the paper defines a custom loss |
+| Training loop | ❌ MINIMAL | Provide a single forward+backward example, not full training |
+| Data pipeline | ❌ SKELETON | Dataset class with TODOs, not full implementation |
+| Evaluation | ⚠️ CONDITIONAL | Metric computation code if paper reports specific metrics |
+| Distributed training | ❌ NEVER | Unless the architecture paper is about distributed/parallel architectures |
+| Visualization | ❌ NEVER | Unless attention visualization IS the contribution |
+
+### Type (b): New training method or loss function
+
+| Component | In scope? | Notes |
+|-----------|-----------|-------|
+| Model architecture | ❌ IMPORT | Import or reference the base architecture used |
+| Loss function | ✅ ALWAYS | This IS (or is part of) the contribution |
+| Training loop | ✅ ALWAYS | The training procedure IS the contribution |
+| Training algorithm | ✅ ALWAYS | Algorithm box should be implemented precisely |
+| Data pipeline | ⚠️ CONDITIONAL | If the training method involves special data handling |
+| Evaluation | ✅ YES | Need to verify the training method produces expected results |
+| LR schedule | ✅ YES | Usually critical for training methods |
+| Optimizer | ✅ YES | If the paper proposes or requires a specific optimizer |
+
+### Type (c): New inference technique
+
+| Component | In scope? | Notes |
+|-----------|-----------|-------|
+| Model architecture | ❌ IMPORT | The model already exists |
+| Inference algorithm | ✅ ALWAYS | This IS the contribution |
+| Training | ❌ NEVER | No training happens at inference time |
+| Data preprocessing | ⚠️ CONDITIONAL | Only if inference requires special preprocessing |
+| Evaluation | ✅ YES | Demonstrate the inference technique produces expected output |
+
+### Type (d): New dataset or benchmark
+
+| Component | In scope? | Notes |
+|-----------|-----------|-------|
+| Dataset class | ✅ ALWAYS | This IS the contribution |
+| Data loading/preprocessing | ✅ ALWAYS | How to load and prepare the data |
+| Evaluation framework | ✅ ALWAYS | Metrics and evaluation protocol |
+| Baseline models | ❌ REFERENCE | Note what baselines were used, don't reimplement |
+| Dataset download | ⚠️ INSTRUCTIONS | Provide clear instructions, don't auto-download |
+
+### Type (e): Theoretical with empirical validation
+
+| Component | In scope? | Notes |
+|-----------|-----------|-------|
+| Experimental setup | ✅ YES | The setup that validates the theory |
+| Measurement code | ✅ YES | What is measured and how |
+| Model architecture | ⚠️ CONDITIONAL | Only if experiments require a specific architecture |
+| Training | ⚠️ CONDITIONAL | Only the specific training setup used in experiments |
+| Proof verification | ❌ NEVER | Proofs are not code |
+
+### Type (f): System/engineering paper
+
+| Component | In scope? | Notes |
+|-----------|-----------|-------|
+| System design | ✅ ALWAYS | This IS the contribution |
+| Performance benchmarks | ⚠️ DOCUMENT | Document the benchmark setup, may not be reproducible |
+| Low-level optimizations | ⚠️ CONDITIONAL | If they can be expressed in Python; CUDA kernels are out of scope |
+| Integration code | ❌ USUALLY NOT | System papers often can't be minimally reproduced |
+
+---
+
+## Mode-specific scope overrides
+
+### minimal mode (default)
+- Stick strictly to the decision tree above
+- Training loop: only if the contribution IS training
+- Data pipeline: skeleton only
+- No distributed training, no checkpointing, no logging
+
+### full mode
+- Same as minimal PLUS:
+- Full training loop with all paper-specified hyperparameters
+- Data pipeline with preprocessing (still needs user to provide data)
+- Evaluation pipeline (not just metric functions)
+- Checkpointing (minimal, for training)
+- Learning rate schedule visualization
+
+### educational mode
+- Same as minimal PLUS:
+- Extra comments explaining ML concepts
+- Extended walkthrough notebook with theory sections
+- PAPER_GUIDE.md that walks through the paper section by section
+
+---
+
+## Explicit exclusions (never in scope regardless of mode)
+
+1. **Baselines and comparison methods.** The paper compares against X, Y, Z — we only implement the paper's contribution, not X, Y, Z.
+
+2. **Ablation variants.** If the paper ablates components A, B, C — we implement the full model with A+B+C, not each variant separately.
+
+3. **Multi-GPU/distributed training infrastructure.** Unless the paper's contribution IS the distributed strategy.
+
+4. **Experiment tracking and logging.** No wandb, no tensorboard, no MLflow setup.
+
+5. **Docker/container configuration.** Not relevant to understanding the paper.
+
+6. **CI/CD, linting, formatting configuration.** This is a research scaffold, not a production project.
+
+7. **Unit tests.** The walkthrough notebook serves as a verification layer. No pytest needed.
+
+8. **Dataset downloading.** Provide instructions, never download automatically.
+
+9. **Pre-trained weight downloading.** Reference where to get them, never download.
+
+10. **Web interfaces, APIs, or demos.** Out of scope entirely.
+
+---
+
+## How to handle scope creep
+
+During code generation, you may be tempted to add "just one more thing" that would make the implementation more complete. Resist this unless:
+
+1. The "one more thing" is directly specified in the paper AND
+2. It is necessary for the core contribution to make sense AND
+3. Leaving it out would make the implementation misleading
+
+If it's "nice to have" — leave it out. A focused, correct implementation of the core contribution is infinitely better than a sprawling implementation with guessed peripherals.
+
+---
+
+## How to document scope decisions
+
+In `REPRODUCTION_NOTES.md`, include a section:
+
+```markdown
+## Scope decisions
+
+### Implemented
+- {Component} — reason: {core contribution / necessary for X / etc.}
+
+### Intentionally excluded
+- {Component} — reason: {baseline method / standard component imported from X / not described in paper}
+
+### Would need for full reproduction
+- {Component} — what it is, where to get it, why we didn't include it
+```
+
+This makes scope decisions transparent. A researcher reading REPRODUCTION_NOTES.md should understand exactly what this implementation covers and what it doesn't.

plugins/codex/plugins/paper2code/skills/paper2code/knowledge/loss_functions.md

@@ -0,0 +1,280 @@
+# Knowledge: Loss Functions
+
+## Purpose
+Canonical implementations of loss functions commonly referenced in papers, with attention to numerical stability, framework-specific differences, and common implementation mistakes.
+
+---
+
+## Cross-Entropy Loss
+
+### Standard cross-entropy (classification)
+
+```python
+# PyTorch built-in:
+loss_fn = nn.CrossEntropyLoss()
+# Expects: predictions (batch, num_classes) as LOGITS (not probabilities)
+#          targets (batch,) as class indices (not one-hot)
+```
+
+**Common mistake:** Passing probabilities (after softmax) instead of logits. PyTorch's `CrossEntropyLoss` applies log-softmax internally for numerical stability.
+
+### Cross-entropy with label smoothing
+
+Papers describe label smoothing differently from what PyTorch does:
+
+**Paper formula (Szegedy et al., 2016):**
+For C classes and smoothing ε:
+- Target for correct class: `1 - ε`
+- Target for each incorrect class: `ε / (C - 1)`
+
+**PyTorch >= 1.10 built-in:**
+```python
+loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)
+```
+PyTorch's implementation distributes ε uniformly: correct class gets `1 - ε + ε/C`, others get `ε/C`. This is slightly different from the paper formula (the paper distributes to incorrect classes only).
+
+**Custom implementation matching the paper exactly:**
+```python
+def label_smoothed_cross_entropy(logits: torch.Tensor, targets: torch.Tensor,
+                                  n_classes: int, smoothing: float = 0.1) -> torch.Tensor:
+    """Label smoothing exactly as described in Szegedy et al., 2016.
+    
+    Args:
+        logits: (batch, n_classes) — raw logits, NOT probabilities
+        targets: (batch,) — class indices
+        n_classes: number of classes
+        smoothing: label smoothing factor ε
+    """
+    log_probs = F.log_softmax(logits, dim=-1)  # (batch, n_classes)
+    
+    # NLL loss for the true class
+    nll_loss = -log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)
+    
+    # Smooth loss: uniform over all classes
+    smooth_loss = -log_probs.mean(dim=-1)
+    
+    loss = (1.0 - smoothing) * nll_loss + smoothing * smooth_loss
+    return loss.mean()
+```
+
+**The difference matters:** For large vocabularies (NLP), the difference between PyTorch's built-in and the paper formula is negligible. For small class counts (e.g., 10), it can be noticeable.
+
+---
+
+## Contrastive Losses
+
+### NT-Xent / InfoNCE (SimCLR, Oord et al.)
+
+```python
+def nt_xent_loss(z_i: torch.Tensor, z_j: torch.Tensor,
+                 temperature: float = 0.5) -> torch.Tensor:
+    """Normalized Temperature-scaled Cross Entropy loss.
+    
+    Used in SimCLR (Chen et al., 2020).
+    
+    Args:
+        z_i: (batch, d) — representations of view 1
+        z_j: (batch, d) — representations of view 2 (positive pairs)
+        temperature: τ — scaling factor
+        
+    Returns:
+        Scalar loss
+    """
+    batch_size = z_i.size(0)
+    
+    # Normalize representations
+    z_i = F.normalize(z_i, dim=-1)  # (batch, d)
+    z_j = F.normalize(z_j, dim=-1)  # (batch, d)
+    
+    # Concatenate representations: [z_i; z_j]
+    z = torch.cat([z_i, z_j], dim=0)  # (2*batch, d)
+    
+    # Compute similarity matrix
+    sim = torch.matmul(z, z.T) / temperature  # (2*batch, 2*batch)
+    
+    # Mask out self-similarity (diagonal)
+    mask = ~torch.eye(2 * batch_size, dtype=torch.bool, device=z.device)
+    sim = sim.masked_fill(~mask, float('-inf'))
+    
+    # Positive pairs: (i, i+batch) and (i+batch, i)
+    labels = torch.cat([
+        torch.arange(batch_size, 2 * batch_size),
+        torch.arange(0, batch_size)
+    ], dim=0).to(z.device)  # (2*batch,)
+    
+    return F.cross_entropy(sim, labels)
+```
+
+**Critical: temperature matters enormously.** SimCLR uses τ=0.5 in the paper, but the appendix shows results are sensitive to this value. CLIP uses τ as a learned parameter initialized to 0.07. Always check what temperature the paper uses.
+
+**Common mistake:** Not normalizing the representations before computing similarity. Without normalization, the cosine similarity becomes a dot product, which is unbounded and destabilizes training.
+
+### Triplet loss
+
+```python
+def triplet_loss(anchor: torch.Tensor, positive: torch.Tensor,
+                 negative: torch.Tensor, margin: float = 1.0) -> torch.Tensor:
+    """Standard triplet margin loss.
+    
+    Args:
+        anchor: (batch, d)
+        positive: (batch, d)
+        negative: (batch, d)
+        margin: minimum desired distance gap
+    """
+    pos_dist = F.pairwise_distance(anchor, positive)  # (batch,)
+    neg_dist = F.pairwise_distance(anchor, negative)  # (batch,)
+    loss = F.relu(pos_dist - neg_dist + margin)
+    return loss.mean()
+```
+
+---
+
+## Diffusion Losses
+
+### DDPM loss (Ho et al., 2020)
+
+The simplified loss from DDPM. This is L_simple from Eq. 14:
+
+```python
+def ddpm_loss(model: nn.Module, x_0: torch.Tensor,
+              noise_schedule: dict) -> torch.Tensor:
+    """Denoising Diffusion Probabilistic Models simplified loss.
+    
+    L_simple = E_{t, x_0, ε}[||ε - ε_θ(x_t, t)||²]
+    
+    The model predicts the noise ε that was added, not the clean signal.
+    
+    Args:
+        model: noise prediction network ε_θ
+        x_0: (batch, C, H, W) — clean images
+        noise_schedule: dict with 'betas', 'alphas_cumprod', etc.
+    """
+    batch_size = x_0.size(0)
+    
+    # Sample random timesteps
+    t = torch.randint(0, len(noise_schedule['betas']), (batch_size,),
+                      device=x_0.device)
+    
+    # Sample noise
+    noise = torch.randn_like(x_0)
+    
+    # Get noisy image: x_t = sqrt(α̅_t) * x_0 + sqrt(1 - α̅_t) * ε
+    alpha_cumprod_t = noise_schedule['alphas_cumprod'][t]  # (batch,)
+    alpha_cumprod_t = alpha_cumprod_t.view(-1, 1, 1, 1)  # (batch, 1, 1, 1) for broadcasting
+    
+    x_t = torch.sqrt(alpha_cumprod_t) * x_0 + torch.sqrt(1 - alpha_cumprod_t) * noise
+    
+    # Predict noise
+    predicted_noise = model(x_t, t)  # (batch, C, H, W)
+    
+    # MSE loss between true and predicted noise
+    loss = F.mse_loss(predicted_noise, noise)
+    return loss
+```
+
+**Noise schedule:**
+```python
+def linear_noise_schedule(timesteps: int, beta_start: float = 1e-4,
+                           beta_end: float = 0.02) -> dict:
+    """Linear variance schedule from DDPM.
+    
+    β_t increases linearly from β_start to β_end over T timesteps.
+    """
+    betas = torch.linspace(beta_start, beta_end, timesteps)
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    
+    return {
+        'betas': betas,
+        'alphas': alphas,
+        'alphas_cumprod': alphas_cumprod,
+        'sqrt_alphas_cumprod': torch.sqrt(alphas_cumprod),
+        'sqrt_one_minus_alphas_cumprod': torch.sqrt(1, - alphas_cumprod),
+    }
+```
+
+**Common mistakes with diffusion losses:**
+1. Indexing: `alphas_cumprod[t]` — make sure `t` is the right shape for broadcasting
+2. The DDPM paper uses 1000 timesteps with linear schedule β₁=1e-4 to β_T=0.02
+3. The model predicts noise ε, not x₀ (in the simplified loss; other parameterizations exist)
+4. `sqrt_one_minus_alphas_cumprod` should NOT go to exactly 0 — check boundary conditions
+
+---
+
+## VAE ELBO
+
+### Evidence Lower Bound
+
+```python
+def vae_loss(recon_x: torch.Tensor, x: torch.Tensor,
+             mu: torch.Tensor, log_var: torch.Tensor,
+             beta: float = 1.0) -> tuple:
+    """VAE loss = Reconstruction loss + β * KL divergence.
+    
+    ELBO = E_q[log p(x|z)] - β * KL(q(z|x) || p(z))
+    
+    Args:
+        recon_x: (batch, ...) — reconstructed output
+        x: (batch, ...) — original input
+        mu: (batch, d_latent) — mean of q(z|x)
+        log_var: (batch, d_latent) — log variance of q(z|x)
+        beta: weight for KL term (β=1 gives standard VAE)
+    """
+    # Reconstruction loss (pixel-wise)
+    recon_loss = F.mse_loss(recon_x, x, reduction='sum') / x.size(0)
+    # Alternative: F.binary_cross_entropy for images in [0,1]
+    
+    # KL divergence: KL(N(μ, σ²) || N(0, 1))
+    # = -0.5 * Σ(1 + log(σ²) - μ² - σ²)
+    kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) / x.size(0)
+    
+    total_loss = recon_loss + beta * kl_loss
+    return total_loss, recon_loss, kl_loss
+```
+
+**Common mistakes:**
+1. **Reconstruction loss type:** Binary cross-entropy for images normalized to [0,1], MSE for others. Papers often don't specify.
+2. **Reduction:** `sum` vs `mean` — affects the relative weight of reconstruction vs KL. If using `mean`, the KL term effectively gets more weight relative to reconstruction for high-dimensional data.
+3. **β value:** β=1 is standard VAE. β-VAE uses β>1 for more disentangled representations. If the paper says "VAE loss" without specifying β, use β=1 and flag it.
+4. **KL divergence sign:** The formula gives a positive KL value. The loss ADDS KL (penalizes deviation from prior). If your KL is negative, the sign is wrong.
+
+---
+
+## Numerical stability patterns
+
+### Log-sum-exp trick
+When computing `log(Σ exp(x_i))`:
+```python
+# WRONG (overflow for large values):
+result = torch.log(torch.sum(torch.exp(x)))
+
+# CORRECT:
+result = torch.logsumexp(x, dim=-1)
+# Internally: max_x + log(Σ exp(x_i - max_x))
+```
+
+### Softmax stability
+PyTorch's `F.softmax` and `F.log_softmax` are already numerically stable (subtract max internally). But if you implement softmax manually:
+```python
+# WRONG:
+weights = torch.exp(scores) / torch.exp(scores).sum(dim=-1, keepdim=True)
+
+# CORRECT:
+scores = scores - scores.max(dim=-1, keepdim=True).values
+weights = torch.exp(scores) / torch.exp(scores).sum(dim=-1, keepdim=True)
+```
+
+### Epsilon in denominators
+When dividing by a value that could be zero (e.g., normalizing):
+```python
+# Not just any epsilon — match the paper's convention or use a safe default
+x_normalized = x / (x.norm(dim=-1, keepdim=True) + 1e-8)
+```
+
+### Clamping log probabilities
+```python
+# When computing log of probabilities that might be 0:
+log_probs = torch.log(probs.clamp(min=1e-8))
+# Or better: compute in log space from the start using log_softmax
+```

plugins/codex/plugins/paper2code/skills/paper2code/knowledge/paper_to_code_mistakes.md

@@ -0,0 +1,332 @@
+# Knowledge: Paper-to-Code Translation Mistakes
+
+## Purpose
+A detailed catalog of patterns where naive paper-to-code translation produces wrong results. These are not bugs in the usual sense — the code runs fine but implements something different from what the paper describes. This file exists because these mistakes are systematic and predictable.
+
+---
+
+## Notation mismatches
+
+### Batch normalization momentum
+
+**The problem:** PyTorch and TensorFlow define momentum differently.
+
+Paper says: `running_mean = (1 - α) * running_mean + α * batch_mean` where α (or momentum) = 0.1
+
+- **PyTorch** `nn.BatchNorm2d(momentum=0.1)`: uses `running_mean = (1 - momentum) * running_mean + momentum * batch_mean` → matches the formula above with momentum=0.1
+- **TensorFlow** `tf.keras.layers.BatchNormalization(momentum=0.99)`: uses `running_mean = momentum * running_mean + (1 - momentum) * batch_mean` → to get the same behavior, set momentum=0.9
+
+**Translation rule:** PyTorch `momentum=x` ≈ TensorFlow `momentum=1-x`
+
+If a paper reports momentum without specifying the framework convention → PARTIALLY_SPECIFIED. Check which framework the official code uses.
+
+### Dropout rate vs keep probability
+
+**The problem:** Older papers (and TensorFlow v1) use "keep probability." Newer papers and PyTorch use "drop probability."
+
+- Paper says "dropout 0.1" → almost certainly means drop probability = 0.1 (keep 90% of neurons)
+- Paper says "dropout rate 0.9" → probably means keep probability = 0.9 (same thing)
+- Paper says "we keep 90% of neurons" → keep probability = 0.9
+
+```python
+# PyTorch: p = drop probability
+nn.Dropout(p=0.1)  # drops 10%, keeps 90%
+```
+
+**When in doubt:** Check the dropout paper (Srivastava et al., 2014) which uses keep probability notation. Most post-2018 papers use drop probability.
+
+### Convolution padding
+
+**The problem:** "same padding" means different things.
+
+- Paper says "same convolution" → output size = input size
+- PyTorch: `nn.Conv2d(padding='same')` or manually compute `padding = kernel_size // 2`
+- TensorFlow: `tf.keras.layers.Conv2D(padding='same')` handles it automatically
+- For even kernel sizes, "same" can't be symmetric → PyTorch requires explicit asymmetric padding via `nn.functional.pad`
+
+### Tensor dimension ordering
+
+- PyTorch images: `(batch, channels, height, width)` — NCHW
+- TensorFlow images: `(batch, height, width, channels)` — NHWC (default)
+- If converting a TensorFlow paper to PyTorch, every convolution, pooling, normalization, and reshape must account for this
+
+---
+
+## Activation function gotchas
+
+### GELU approximation
+
+There are two common GELU implementations:
+
+```python
+# Exact GELU (PyTorch default since 1.12):
+nn.GELU()
+# = x * Φ(x) where Φ is the standard normal CDF
+
+# Tanh approximation (used in GPT-2, BERT):
+nn.GELU(approximate='tanh')
+# = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))
+```
+
+These give slightly different results. Some papers use one, some the other. BERT uses the tanh approximation. Modern papers usually use exact GELU but often don't specify.
+
+### Swish / SiLU
+
+```python
+# These are the same thing:
+nn.SiLU()  # PyTorch name
+# = x * sigmoid(x)
+
+# Some papers call it Swish with a trainable β:
+# swish(x) = x * sigmoid(β * x)
+# When β=1, this is SiLU. PyTorch's SiLU always uses β=1.
+```
+
+If the paper says "Swish" without specifying β, use β=1 (SiLU) and flag it.
+
+### Leaky ReLU slope
+
+Papers rarely specify the negative slope for Leaky ReLU:
+- Paper default: 0.01
+- PyTorch default: 0.01 (matches)
+- But some implementations use 0.1 or 0.2 (especially in GANs)
+
+---
+
+## Weight initialization gotchas
+
+### Xavier/Glorot initialization
+
+```python
+# Xavier uniform (Glorot & Bengio, 2010):
+nn.init.xavier_uniform_(layer.weight)  # U(-a, a) where a = sqrt(6 / (fan_in + fan_out))
+
+# Xavier normal:
+nn.init.xavier_normal_(layer.weight)  # N(0, 2 / (fan_in + fan_out))
+```
+
+**The problem:** "Xavier initialization" doesn't specify uniform or normal. They're different. Papers rarely specify which one.
+
+### Kaiming/He initialization
+
+```python
+# Kaiming uniform (He et al., 2015):
+nn.init.kaiming_uniform_(layer.weight, mode='fan_in', nonlinearity='relu')
+
+# Kaiming normal:
+nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')
+```
+
+**The problem:** `mode='fan_in'` vs `mode='fan_out'` matters. PyTorch defaults to `fan_in`, but some papers require `fan_out` for the decoder.
+
+### PyTorch default initialization
+
+If the paper doesn't mention initialization, PyTorch applies:
+- `nn.Linear`: Kaiming uniform (fan_in)
+- `nn.Conv2d`: Kaiming uniform (fan_in)
+- `nn.Embedding`: Normal(0, 1)
+- `nn.LayerNorm`: weight=1, bias=0
+- `nn.BatchNorm`: weight=1, bias=0
+
+These defaults are reasonable but may not match what the paper used (and didn't specify).
+
+---
+
+## Dimension and reshaping gotchas
+
+### view() vs reshape()
+
+```python
+# view() requires contiguous memory — will error if not contiguous
+x = x.view(batch, -1)
+
+# reshape() handles non-contiguous tensors by copying if necessary
+x = x.reshape(batch, -1)
+```
+
+After `transpose()` or `permute()`, tensors are NOT contiguous. Use `.contiguous()` before `.view()` or just use `.reshape()`.
+
+### einsum notation
+
+Papers increasingly use einsum notation. Common patterns:
+
+```python
+# Matrix multiplication:
+torch.einsum('bij,bjk->bik', A, B)  # same as torch.bmm(A, B)
+
+# Attention scores:
+torch.einsum('bhqd,bhkd->bhqk', q, k)  # (batch, heads, q_len, k_len)
+
+# Bilinear:
+torch.einsum('bi,ij,bj->b', x, W, y)
+```
+
+**Gotcha:** einsum is correct but can be slow for some patterns. For standard operations (matmul, dot product), use `torch.matmul` or `@` operator — they're optimized.
+
+---
+
+## Loss function gotchas
+
+### Cross-entropy: logits vs probabilities
+
+If you see `loss = -sum(y * log(p))` in a paper, this is cross-entropy applied to probabilities.
+PyTorch `nn.CrossEntropyLoss` expects LOGITS and applies log-softmax internally.
+
+```python
+# WRONG: passing softmax output to CrossEntropyLoss
+probs = F.softmax(logits, dim=-1)
+loss = nn.CrossEntropyLoss()(probs, targets)  # WRONG — double softmax
+
+# CORRECT:
+loss = nn.CrossEntropyLoss()(logits, targets)  # applies log-softmax internally
+```
+
+### MSE reduction
+
+Papers often specify loss as a sum or average differently:
+- PyTorch `MSELoss(reduction='mean')`: averages over ALL elements (batch × features)
+- Some papers average only over features, then sum over batch
+- Some papers sum everything, no averaging
+
+Check the paper's equation: does it have $\frac{1}{N}$ (mean) or $\sum$ (sum)?
+
+### KL divergence direction
+
+$KL(P || Q) \neq KL(Q || P)$
+
+- Forward KL ($KL(P || Q)$): "I want Q to cover everywhere P has mass" — mode-covering
+- Reverse KL ($KL(Q || P)$): "I want Q to be 0 wherever P is 0" — mode-seeking
+
+PyTorch's `nn.KLDivLoss` computes `KL(target || input)` — note the argument order!
+
+```python
+# PyTorch expects LOG probabilities as input, regular probabilities as target
+loss = nn.KLDivLoss(reduction='batchmean')(
+    F.log_softmax(predicted_logits, dim=-1),  # LOG probs
+    target_probs  # regular probs
+)
+```
+
+---
+
+## Training loop gotchas
+
+### optimizer.zero_grad() placement
+
+```python
+# CORRECT: zero gradients BEFORE forward pass
+optimizer.zero_grad()
+output = model(input)
+loss = loss_fn(output, target)
+loss.backward()
+optimizer.step()
+
+# ALSO CORRECT (set_to_none is faster):
+optimizer.zero_grad(set_to_none=True)
+```
+
+### Learning rate scheduler step
+
+Different schedulers step at different rates:
+```python
+# Per-step schedulers (most warmup schedules):
+for step in range(total_steps):
+    train_step()
+    scheduler.step()  # after every training step
+
+# Per-epoch schedulers (ReduceLROnPlateau, StepLR):
+for epoch in range(epochs):
+    train_epoch()
+    scheduler.step()  # after every epoch
+```
+
+Papers often don't specify when the scheduler steps. If it's a warmup schedule, it's per-step.
+
+### Gradient accumulation
+
+```python
+# If effective_batch = micro_batch × accumulation_steps:
+for i, batch in enumerate(dataloader):
+    loss = model(batch) / accumulation_steps  # DIVIDE by accumulation steps
+    loss.backward()  # accumulate gradients
+    
+    if (i + 1) % accumulation_steps == 0:
+        optimizer.step()
+        optimizer.zero_grad()
+```
+
+**Common mistake:** Forgetting to divide the loss by `accumulation_steps`. Without this, the effective learning rate is scaled by the number of accumulation steps.
+
+---
+
+## Evaluation gotchas
+
+### BLEU score implementations
+
+Different BLEU implementations give different numbers:
+- `sacrebleu` (recommended, reproducible): uses standard tokenization
+- `nltk.translate.bleu_score`: requires pre-tokenized input
+- `torchtext.data.metrics.bleu_score`: deprecated
+- Papers before 2020 may use any of these
+
+**Difference can be 1-2 BLEU points** — significant for reporting.
+
+### FID (Fréchet Inception Distance)
+
+- Must use the SAME Inception v3 model weights (the default TF weights)
+- PyTorch and TensorFlow Inception models give slightly different features
+- Number of samples used for FID computation matters (more ≈ more stable)
+- `clean-fid` package is the recommended implementation
+
+### Top-k accuracy
+
+```python
+def topk_accuracy(output: torch.Tensor, target: torch.Tensor, k: int = 5) -> float:
+    """Compute top-k accuracy.
+    output: (batch, num_classes) — logits
+    target: (batch,) — class indices
+    """
+    _, pred = output.topk(k, dim=-1)  # (batch, k)
+    correct = pred.eq(target.unsqueeze(-1)).any(dim=-1)  # (batch,)
+    return correct.float().mean().item()
+```
+
+---
+
+## Miscellaneous gotchas
+
+### model.eval() vs model.train()
+
+```python
+model.eval()   # Disables dropout and uses running stats for batch norm
+model.train()  # Enables dropout and computes batch stats for batch norm
+
+# CRITICAL: Always call model.eval() before evaluation/inference
+# Forgetting this is a very common source of poor eval performance
+```
+
+### torch.no_grad() vs torch.inference_mode()
+
+```python
+# For evaluation (no gradient computation):
+with torch.no_grad():
+    output = model(input)
+    
+# For inference (even faster, more restricted):
+with torch.inference_mode():
+    output = model(input)
+```
+
+### Determinism
+
+If the paper reports "we average over 3 random seeds":
+```python
+def set_seed(seed: int):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+```
+Note: `cudnn.deterministic = True` can significantly slow down training.

plugins/codex/plugins/paper2code/skills/paper2code/knowledge/training_recipes.md

@@ -0,0 +1,253 @@
+# Knowledge: Training Recipes
+
+## Purpose
+Common optimizer configurations, learning rate schedules, and training details that papers frequently assume you know but don't re-explain. When a paper says "we use standard training settings," this file helps you determine what they probably mean (and helps you flag what's still ambiguous).
+
+---
+
+## Optimizers
+
+### Adam (Kingma & Ba, 2014)
+
+Default parameters in the paper: β₁=0.9, β₂=0.999, ε=1e-8
+
+**PyTorch defaults match the Adam paper.**
+
+But in practice:
+- Transformer papers often use β₂=0.98 (from Vaswani et al.)
+- Some NLP papers use β₂=0.95
+- ε values vary: 1e-8 (default), 1e-6 (BERT), 1e-9 (Vaswani et al.)
+- "We use Adam" does NOT mean these defaults — you must check
+
+### AdamW (Loshchilov & Hutter, 2019)
+
+AdamW is NOT the same as Adam with L2 regularization:
+- Adam + L2: weight decay is applied to the gradient, which scales with the adaptive learning rate
+- AdamW: weight decay is applied directly to the weights, decoupled from the gradient update
+- This difference matters for large learning rates
+
+```python
+# Adam with L2 (WRONG if paper says "AdamW" or "decoupled weight decay"):
+optimizer = torch.optim.Adam(params, lr=lr, weight_decay=0.01)
+
+# AdamW (correct decoupled weight decay):
+optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=0.01)
+```
+
+Critical detail: **Weight decay is usually NOT applied to bias terms and LayerNorm parameters.** If the paper doesn't state this, implement it anyway and flag as [ASSUMPTION]:
+
+```python
+# [ASSUMPTION] Not applying weight decay to biases and normalization layers
+# (standard practice, but paper does not specify)
+no_decay = ['bias', 'LayerNorm.weight', 'LayerNorm.bias']
+param_groups = [
+    {'params': [p for n, p in model.named_parameters()
+                if not any(nd in n for nd in no_decay)],
+     'weight_decay': config.weight_decay},
+    {'params': [p for n, p in model.named_parameters()
+                if any(nd in n for nd in no_decay)],
+     'weight_decay': 0.0}
+]
+optimizer = torch.optim.AdamW(param_groups, lr=config.lr)
+```
+
+### SGD with momentum
+
+Common in vision papers (ResNet, EfficientNet):
+- Momentum: 0.9 (almost always)
+- Weight decay: 1e-4 (vision standard, not universal)
+- Nesterov: sometimes on, sometimes off — papers often don't specify
+
+---
+
+## Learning rate schedules
+
+### Linear warmup then constant
+```python
+def get_warmup_schedule(optimizer, warmup_steps):
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return step / warmup_steps
+        return 1.0
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+```
+
+### Linear warmup then linear decay
+```python
+def get_linear_schedule(optimizer, warmup_steps, total_steps):
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return step / warmup_steps
+        return max(0.0, (total_steps - step) / (total_steps - warmup_steps))
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+```
+
+### Linear warmup then cosine decay (very common)
+```python
+def get_cosine_schedule(optimizer, warmup_steps, total_steps, min_lr_ratio=0.0):
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return step / warmup_steps
+        progress = (step - warmup_steps) / (total_steps - warmup_steps)
+        return min_lr_ratio + (1.0 - min_lr_ratio) * 0.5 * (1.0 + math.cos(math.pi * progress))
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+```
+
+### Transformer schedule (Vaswani et al., 2017)
+
+The original transformer uses a unique schedule that most papers reference:
+```python
+def get_transformer_schedule(optimizer, d_model, warmup_steps):
+    """lr = d_model^(-0.5) * min(step^(-0.5), step * warmup_steps^(-1.5))"""
+    def lr_lambda(step):
+        step = max(step, 1)  # avoid division by zero
+        return d_model ** (-0.5) * min(step ** (-0.5), step * warmup_steps ** (-1.5))
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+```
+
+Note: This schedule doesn't need a base learning rate — it computes the LR from scratch.
+When papers say "following Vaswani et al. for learning rate," this is what they mean.
+
+---
+
+## Batch size semantics
+
+### The nomenclature problem
+
+When a paper says "batch size 256," it could mean:
+1. **Per-GPU batch size**: each GPU processes 256 samples
+2. **Global batch size**: 256 total across all GPUs
+3. **Effective batch size after gradient accumulation**: might be 32 per-GPU × 8 accumulation steps
+
+This is CRITICAL for reproduction because learning rate scaling depends on batch size.
+
+### How to determine which
+
+Check for:
+- "per-GPU batch size" or "micro-batch size" → clear
+- "total batch size" or "effective batch size" → clear
+- Just "batch size" → ambiguous. Look for:
+  - Number of GPUs mentioned → if N GPUs and batch size B, probably B total
+  - "gradient accumulation steps = K" → if batch B and K steps, effective batch = B × K
+  - If they mention both batch size and number of GPUs but don't clarify → flag as PARTIALLY_SPECIFIED
+
+### Learning rate linear scaling rule
+
+Many papers implicitly use the linear scaling rule (Goyal et al., 2017):
+- If base LR is `lr` for batch size `B`, then LR for batch size `B'` is `lr * B' / B`
+- This applies to SGD. For Adam, the scaling is usually sqrt: `lr * sqrt(B' / B)`
+- If the paper trains on 8 GPUs with batch 256 and you use 1 GPU with batch 32, adjust the LR
+- **Flag this if you change the batch size from what the paper specifies**
+
+---
+
+## Gradient clipping
+
+### Max gradient norm (most common)
+```python
+torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+```
+Apply AFTER `loss.backward()` and BEFORE `optimizer.step()`.
+
+Common values:
+- 1.0 (most common)
+- 0.5 (some NLP papers)
+- 5.0 (some RL papers)
+- If paper says "gradient clipping" without a value → UNSPECIFIED
+
+### Max gradient value (less common)
+```python
+torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)
+```
+Different from norm clipping — clips each gradient element independently.
+
+---
+
+## Mixed precision training
+
+### What it means
+Train with FP16 (or BF16) for speed, but keep master weights in FP32 for stability.
+
+### PyTorch implementation
+```python
+scaler = torch.amp.GradScaler()
+with torch.amp.autocast(device_type='cuda', dtype=torch.float16):
+    output = model(input)
+    loss = loss_fn(output, target)
+
+scaler.scale(loss).backward()
+scaler.step(optimizer)
+scaler.update()
+```
+
+### What papers don't tell you
+- BF16 doesn't need loss scaling (GradScaler) — it has the same dynamic range as FP32
+- FP16 DOES need loss scaling to prevent gradient underflow
+- Some operations must stay in FP32: softmax, layer norm, loss computation
+- PyTorch autocast handles most of this automatically, but custom operations may not be covered
+- "We use mixed precision" without specifying FP16 vs BF16 is **PARTIALLY_SPECIFIED**
+
+---
+
+## Exponential Moving Average (EMA)
+
+Common in diffusion models, GANs, and some vision models. Papers often mention EMA without specifying the decay rate.
+
+```python
+class EMA:
+    """Maintains exponential moving average of model parameters."""
+    def __init__(self, model: nn.Module, decay: float = 0.9999):
+        self.decay = decay
+        self.shadow = {name: param.clone().detach()
+                      for name, param in model.named_parameters()}
+
+    @torch.no_grad()
+    def update(self, model: nn.Module):
+        for name, param in model.named_parameters():
+            self.shadow[name].mul_(self.decay).add_(param.data, alpha=1 - self.decay)
+
+    def apply(self, model: nn.Module):
+        """Load EMA weights into model for evaluation."""
+        for name, param in model.named_parameters():
+            param.data.copy_(self.shadow[name])
+```
+
+Common decay values:
+- 0.9999 (diffusion models — DDPM, DDIM)
+- 0.999 (some GAN papers)
+- 0.99 (faster averaging)
+- If paper says "EMA" but not the decay rate → UNSPECIFIED
+
+---
+
+## Common training recipes by domain
+
+### Language models (GPT-style)
+- Optimizer: AdamW with β₁=0.9, β₂=0.95, ε=1e-8
+- Weight decay: 0.1 (not on biases/norms)
+- LR schedule: cosine decay with linear warmup
+- Gradient clipping: 1.0 (max norm)
+- Batch size: typically reported as total tokens/batch (e.g., "batch of 0.5M tokens")
+
+### Vision transformers (ViT-style)
+- Optimizer: AdamW with β₁=0.9, β₂=0.999
+- Weight decay: 0.05-0.3
+- LR schedule: cosine decay with linear warmup (5-10 epochs warmup)
+- Data augmentation: RandAugment, Mixup, CutMix, random erasing (but specific combination varies)
+- Label smoothing: 0.1
+
+### Diffusion models (DDPM-style)
+- Optimizer: Adam with β₁=0.9, β₂=0.999 (or 0.9999)
+- LR: constant (often 2e-4 or 1e-4)
+- No LR warmup (sometimes)
+- EMA: 0.9999
+- Gradient clipping: sometimes but not always
+
+### Contrastive learning (SimCLR-style)
+- Optimizer: SGD with momentum 0.9 (or LARS)
+- LR schedule: cosine decay after linear warmup (10 epochs)
+- Weight decay: 1e-6 (very small)
+- Batch size: LARGE (4096+) — crucial for performance
+- Temperature: 0.1 or 0.5
+
+**WARNING:** These are guidelines, not specifications. If a paper says "we use standard training," it MIGHT mean the recipe above — but it's still UNSPECIFIED unless the paper states specifics. Use this knowledge to choose reasonable defaults, not to skip flagging.

plugins/codex/plugins/paper2code/skills/paper2code/knowledge/transformer_components.md

@@ -0,0 +1,362 @@
+# Knowledge: Transformer Components
+
+## Purpose
+Canonical correct implementations of transformer building blocks that papers frequently reference but don't re-explain. When a paper says "standard transformer encoder," this file tells you what that means and what mistakes to avoid.
+
+---
+
+## Multi-Head Attention
+
+### Canonical implementation
+
+```python
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0,
+                 bias: bool = True):
+        super().__init__()
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_k = d_model // n_heads  # head dimension
+
+        self.W_q = nn.Linear(d_model, d_model, bias=bias)
+        self.W_k = nn.Linear(d_model, d_model, bias=bias)
+        self.W_v = nn.Linear(d_model, d_model, bias=bias)
+        self.W_o = nn.Linear(d_model, d_model, bias=bias)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
+                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        batch_size = query.size(0)
+
+        # Project and reshape: (batch, seq, d_model) -> (batch, n_heads, seq, d_k)
+        q = self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+        k = self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+        v = self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+
+        # Scaled dot-product attention
+        # (batch, n_heads, seq_q, d_k) @ (batch, n_heads, d_k, seq_k) -> (batch, n_heads, seq_q, seq_k)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, float('-inf'))
+
+        attn_weights = F.softmax(scores, dim=-1)  # (batch, n_heads, seq_q, seq_k)
+        attn_weights = self.dropout(attn_weights)
+
+        # (batch, n_heads, seq_q, seq_k) @ (batch, n_heads, seq_k, d_k) -> (batch, n_heads, seq_q, d_k)
+        context = torch.matmul(attn_weights, v)
+
+        # Reshape back: (batch, n_heads, seq_q, d_k) -> (batch, seq_q, d_model)
+        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
+        return self.W_o(context)  # (batch, seq_q, d_model)
+```
+
+### Common mistakes
+
+1. **Scaling by sqrt(d_model) instead of sqrt(d_k)**
+   - The scale factor is `sqrt(d_k)` where `d_k = d_model / n_heads`
+   - NOT `sqrt(d_model)`. This is the single most common mistake.
+   - Vaswani et al. §3.2.1: "We suspect that for large values of d_k, the dot products grow large in magnitude"
+
+2. **Wrong mask convention**
+   - Additive mask: add a large negative number (e.g., -1e9 or -inf) to scores BEFORE softmax
+   - Multiplicative mask: multiply attention weights by 0/1 AFTER softmax
+   - Both are valid, but papers rarely specify which. Additive is more common and numerically cleaner.
+   - Using -inf is cleaner than -1e9 (avoids non-zero attention for -1e9 with float16)
+
+3. **Forgetting `.contiguous()` after transpose**
+   - After `transpose(1, 2)`, the tensor may not be contiguous
+   - `.view()` requires a contiguous tensor
+   - This will crash, not silently fail — but it's a common "why doesn't my code run" bug
+
+4. **Causal mask shape**
+   - Should be `(1, 1, seq_len, seq_len)` for broadcasting with `(batch, n_heads, seq_len, seq_len)` scores
+   - Mask where `mask[i][j] = 1` if position `j` is allowed for position `i`
+   - Upper triangular = disallowed, not lower triangular (common mistake)
+
+### Causal masking
+
+```python
+def create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
+    """Creates a causal (autoregressive) attention mask.
+    Returns a (1, 1, seq_len, seq_len) boolean tensor where True = attend, False = mask.
+    """
+    mask = torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool))
+    return mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, seq_len)
+```
+
+---
+
+## Positional Encodings
+
+### Sinusoidal (Vaswani et al., 2017)
+
+```python
+class SinusoidalPositionalEncoding(nn.Module):
+    """Fixed sinusoidal positional encoding from 'Attention Is All You Need'.
+
+    PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
+    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
+    """
+    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.0):
+        super().__init__()
+        self.dropout = nn.Dropout(dropout)
+
+        pe = torch.zeros(max_len, d_model)  # (max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1).float()  # (max_len, 1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
+        )  # (d_model/2,)
+
+        pe[:, 0::2] = torch.sin(position * div_term)  # even indices
+        pe[:, 1::2] = torch.cos(position * div_term)  # odd indices
+        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (batch, seq_len, d_model)
+        x = x + self.pe[:, :x.size(1)]
+        return self.dropout(x)
+```
+
+### Common mistakes with sinusoidal PE:
+- Using `arange(0, d_model)` instead of `arange(0, d_model, 2)` for div_term
+- Off-by-one in position indexing (should start at 0)
+- Forgetting to `register_buffer` (so it's not a parameter but moves with the model to GPU)
+
+### Learned positional embeddings
+
+```python
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_len: int, d_model: int):
+        super().__init__()
+        self.embedding = nn.Embedding(max_len, d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (batch, seq_len, d_model)
+        positions = torch.arange(x.size(1), device=x.device)  # (seq_len,)
+        return x + self.embedding(positions)  # broadcast over batch
+```
+
+### Rotary Position Embedding (RoPE) — Su et al., 2021
+
+```python
+class RotaryPositionalEmbedding(nn.Module):
+    """RoPE: Enhanced Transformer with Rotary Position Embedding.
+
+    Applied to each head individually within the attention computation,
+    AFTER the Q and K projections but BEFORE the dot product.
+    """
+    def __init__(self, d_head: int, max_len: int = 8192, base: float = 10000.0):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, d_head, 2).float() / d_head))
+        self.register_buffer('inv_freq', inv_freq)
+        self.max_len = max_len
+
+    def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        t = torch.arange(seq_len, device=x.device).float()
+        freqs = torch.outer(t, self.inv_freq)  # (seq_len, d_head/2)
+        emb = torch.cat([freqs, freqs], dim=-1)  # (seq_len, d_head)
+        return emb.cos(), emb.sin()
+
+
+def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Apply RoPE to query or key tensor.
+    x: (batch, n_heads, seq_len, d_head)
+    """
+    d_half = x.shape[-1] // 2
+    x1, x2 = x[..., :d_half], x[..., d_half:]
+    return torch.cat([
+        x1 * cos[..., :d_half] - x2 * sin[..., :d_half],
+        x2 * cos[..., d_half:] + x1 * sin[..., d_half:]
+    ], dim=-1)
+```
+
+### Key difference: RoPE is applied to Q and K individually, NOT summed onto embeddings like sinusoidal PE.
+
+### ALiBi (Press et al., 2022)
+
+ALiBi doesn't use positional embeddings at all. Instead, it adds a linear bias to the attention scores:
+
+```python
+def get_alibi_slopes(n_heads: int) -> torch.Tensor:
+    """Compute ALiBi slopes for each head.
+    Head i gets slope 2^(-8i/n_heads) for i = 1, ..., n_heads
+    """
+    ratio = 2 ** (-8 / n_heads)
+    slopes = torch.tensor([ratio ** i for i in range(1, n_heads + 1)])
+    return slopes  # (n_heads,)
+
+def apply_alibi(scores: torch.Tensor, slopes: torch.Tensor) -> torch.Tensor:
+    """Apply ALiBi bias to attention scores.
+    scores: (batch, n_heads, seq_q, seq_k)
+    slopes: (n_heads,)
+    """
+    seq_q, seq_k = scores.size(-2), scores.size(-1)
+    # Position difference: relative distance between query and key positions
+    positions = torch.arange(seq_k, device=scores.device).unsqueeze(0) - \
+                torch.arange(seq_q, device=scores.device).unsqueeze(1)  # (seq_q, seq_k)
+    bias = slopes.unsqueeze(-1).unsqueeze(-1) * positions.unsqueeze(0)  # (n_heads, seq_q, seq_k)
+    return scores + bias.unsqueeze(0)  # broadcast over batch
+```
+
+---
+
+## Layer Normalization
+
+### Pre-norm vs Post-norm — THIS MATTERS ENORMOUSLY
+
+**Post-norm (original Transformer):**
+```python
+# Post-norm: normalize AFTER the residual addition
+x = self.norm(x + self.sublayer(x))
+```
+
+**Pre-norm (GPT-2, most modern transformers):**
+```python
+# Pre-norm: normalize BEFORE the sublayer, residual OUTSIDE the norm
+x = x + self.sublayer(self.norm(x))
+```
+
+**Why it matters:**
+- Post-norm requires learning rate warmup and careful initialization
+- Pre-norm is much more stable to train at scale
+- They produce different quality models — not interchangeable
+- Many papers show post-norm in figures but use pre-norm in experiments — ALWAYS CHECK
+
+### RMSNorm (Zhang & Sennrich, 2019)
+
+```python
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization.
+    Used in LLaMA, T5. Simpler than LayerNorm (no centering, no bias).
+    """
+    def __init__(self, d_model: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d_model))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
+        return x / rms * self.weight
+```
+
+---
+
+## Feed-Forward Network
+
+### Standard (Vaswani et al.)
+```python
+class FeedForward(nn.Module):
+    """Two-layer feed-forward network with expansion factor.
+    FFN(x) = W_2 * activation(W_1 * x + b_1) + b_2
+    """
+    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0,
+                 activation: str = "relu"):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, d_ff)
+        self.linear2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+        if activation == "relu":
+            self.activation = nn.ReLU()
+        elif activation == "gelu":
+            self.activation = nn.GELU()
+        elif activation == "silu":
+            self.activation = nn.SiLU()
+        else:
+            raise ValueError(f"Unknown activation: {activation}")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # (batch, seq, d_model) -> (batch, seq, d_ff) -> (batch, seq, d_model)
+        return self.linear2(self.dropout(self.activation(self.linear1(x))))
+```
+
+### SwiGLU (Shazeer, 2020) — used in LLaMA, PaLM
+```python
+class SwiGLU(nn.Module):
+    """Gated feed-forward with SiLU activation.
+    SwiGLU(x) = (SiLU(W_1 * x) ⊙ W_3 * x) * W_2
+    Note: uses 3 weight matrices, not 2. This changes parameter count.
+    """
+    def __init__(self, d_model: int, d_ff: int):
+        super().__init__()
+        self.w1 = nn.Linear(d_model, d_ff, bias=False)
+        self.w2 = nn.Linear(d_ff, d_model, bias=False)
+        self.w3 = nn.Linear(d_model, d_ff, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+```
+
+---
+
+## Embedding with Weight Tying
+
+```python
+class TransformerEmbedding(nn.Module):
+    """Token + positional embedding with optional weight tying to output projection.
+
+    Weight tying (Press & Wolf, 2017): The embedding matrix and the output
+    projection matrix are the SAME tensor. This reduces parameters and often
+    improves performance. Many papers do this without mentioning it explicitly.
+    """
+    def __init__(self, vocab_size: int, d_model: int, max_len: int,
+                 dropout: float = 0.0, scale: bool = True):
+        super().__init__()
+        self.token_emb = nn.Embedding(vocab_size, d_model)
+        self.pos_emb = LearnedPositionalEmbedding(max_len, d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.scale = math.sqrt(d_model) if scale else 1.0
+        # Vaswani et al. §3.4: "we multiply those weights by sqrt(d_model)"
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (batch, seq_len) of token IDs
+        tok = self.token_emb(x) * self.scale  # (batch, seq_len, d_model)
+        return self.dropout(self.pos_emb(tok))
+```
+
+**Weight tying note:** If the paper ties embedding and output weights, the output projection is `F.linear(x, model.embedding.token_emb.weight)` — not a separate `nn.Linear`. Many papers do this without stating it. Check the parameter count in the paper against your model — if yours is higher, weight tying might be missing.
+
+---
+
+## Complete Transformer Block
+
+### Post-norm variant (original)
+```python
+class TransformerBlockPostNorm(nn.Module):
+    def __init__(self, d_model, n_heads, d_ff, dropout):
+        super().__init__()
+        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
+        self.ff = FeedForward(d_model, d_ff, dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+    def forward(self, x, mask=None):
+        x = self.norm1(x + self.dropout1(self.attn(x, x, x, mask)))
+        x = self.norm2(x + self.dropout2(self.ff(x)))
+        return x
+```
+
+### Pre-norm variant (modern standard)
+```python
+class TransformerBlockPreNorm(nn.Module):
+    def __init__(self, d_model, n_heads, d_ff, dropout):
+        super().__init__()
+        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
+        self.ff = FeedForward(d_model, d_ff, dropout)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+    def forward(self, x, mask=None):
+        x = x + self.dropout1(self.attn(self.norm1(x), self.norm1(x), self.norm1(x), mask))
+        x = x + self.dropout2(self.ff(self.norm2(x)))
+        return x
+```
+
+**Key difference:** In pre-norm, LayerNorm is applied BEFORE each sublayer. The residual connection adds the UN-normalized input. This is more stable for training deep models.

plugins/codex/plugins/paper2code/skills/paper2code/pipeline/01_paper_acquisition.md

@@ -0,0 +1,152 @@
+# Stage 1: Paper Acquisition and Parsing
+
+## Purpose
+Fetch the arxiv paper, extract its full text with mathematical notation preserved, and produce a structured markdown representation that downstream stages can consume section by section.
+
+## Input
+- `ARXIV_ID`: e.g., `2106.09685` or `2106.09685v2`
+
+## Output
+- `.paper2code_work/{ARXIV_ID}/paper_text.md` — full paper text in markdown
+- `.paper2code_work/{ARXIV_ID}/paper_metadata.json` — title, authors, abstract, categories
+- `.paper2code_work/{ARXIV_ID}/sections/` — individual section files
+- `.paper2code_work/{ARXIV_ID}/algorithms/` — extracted algorithm boxes
+- `.paper2code_work/{ARXIV_ID}/equations/` — extracted numbered equations
+- `.paper2code_work/{ARXIV_ID}/tables/` — extracted tables (especially hyperparameter tables)
+- `.paper2code_work/{ARXIV_ID}/footnotes.md` — all footnotes collected
+
+---
+
+## Reasoning protocol
+
+### Step 1: Normalize the input
+
+Ask yourself:
+- Is this a full URL or a bare ID?
+- Does it have a version suffix (v1, v2)?
+- Is the ID format valid? (should be YYMM.NNNNN or older format like arch-ive/NNNNNNN)
+
+Strip to just the ID. Keep version suffix if present.
+
+### Step 2: Fetch the paper
+
+Run `scripts/fetch_paper.py`. The script handles:
+1. Downloading from `https://arxiv.org/pdf/{id}.pdf`
+2. Extraction via `pymupdf4llm` (preferred — preserves math notation as LaTeX)
+3. Fallback to `pdfplumber` if pymupdf4llm fails
+4. Fallback to HTML from `https://ar5iv.labs.arxiv.org/html/{id}` if PDF extraction produces garbled output
+
+**How to detect garbled output:** After extraction, scan the first 500 characters. If more than 20% are non-ASCII, non-LaTeX special characters, or if the text has no recognizable English words, the extraction is garbled. Fall back.
+
+### Step 3: Verify extraction quality
+
+Read the extracted `paper_text.md` and check:
+- [ ] Can you identify the paper title?
+- [ ] Is the abstract present and readable?
+- [ ] Are section headings identifiable?
+- [ ] Are equations present (even if in LaTeX notation)?
+- [ ] Is the references section present at the end?
+
+If any of these fail, attempt the ar5iv HTML fallback. If that also fails, inform the user that automatic extraction failed and ask them to paste the paper text directly.
+
+### Step 4: Run structure extraction
+
+Run `scripts/extract_structure.py` on the extracted text. This script:
+- Identifies section boundaries using heading patterns (`#`, numbered headings, ALL CAPS headings)
+- Extracts algorithm boxes (text between "Algorithm N" and the next section)
+- Extracts numbered equations
+- Extracts tables (especially those containing hyperparameters, learning rates, dimensions)
+- Extracts footnotes
+
+### Step 5: Verify all critical sections exist
+
+**You MUST find these sections** (they may be named differently):
+- Abstract
+- Introduction (or "1 Introduction" or similar)
+- Method/Model/Approach section (this is the core — it may be named anything)
+- Experiments/Results
+- Conclusion
+
+**You MUST actively look for these** (authors hide crucial details here):
+- Appendix — check for content AFTER the references. Many papers have appendices with implementation details, hyperparameter tables, ablation studies, prompts, and proofs that are essential for reproduction
+- Supplementary material references — if the paper mentions "see supplementary" or "see appendix," note what is referenced
+- Footnotes — often contain critical caveats about implementation choices
+
+### Step 6: Special handling for appendices
+
+This deserves its own step because it's that important:
+
+1. After the References section, look for any additional content (Appendix A, B, C, etc.)
+2. If appendices exist, extract them as separate section files
+3. Pay special attention to:
+   - Hyperparameter tables (often in Appendix A or B)
+   - System Prompts
+   - Architecture diagrams described in text
+   - Training details (often Appendix C or D)
+   - Ablation studies (contain information about what matters and what doesn't)
+4. If the paper references a supplementary PDF, note this — you may need to fetch it separately from the arxiv page
+
+### Step 7: Extract metadata
+
+From the paper text or arxiv page, extract:
+- Title
+- Authors
+- Year
+- Arxiv categories (e.g., cs.LG, cs.CV)
+- Abstract (first 500 words)
+
+Save to `paper_metadata.json`.
+
+### Step 8: Search for official code repositories
+
+The `fetch_paper.py` script automatically searches for official code in two places:
+
+1. **Inside the paper text** — scans for GitHub/GitLab/Bitbucket URLs and phrases like "code available at," "our implementation is released at," etc.
+2. **The arxiv abstract page** — checks for code repository links in the page HTML.
+
+Results are saved to `paper_metadata.json` under the `official_code` key. Each entry has:
+- `url` — the repository URL
+- `source` — where it was found (`paper_text` or `arxiv_page`)
+- `context` — surrounding text that confirms it's the authors' code
+
+**After the script runs, verify the links:**
+- Open the repository URL. Is it actually the authors' official code for THIS paper, or an unrelated repo?
+- Does the repo contain a working implementation? Some repos are empty placeholders or "coming soon."
+- Note the primary language/framework — this may inform your implementation choices.
+
+If official code is found, it becomes a critical resource for Stage 3 (Ambiguity Audit). Every `[UNSPECIFIED]` item should be checked against the official code before choosing a default. Choices resolved this way get the `[FROM_OFFICIAL_CODE]` tag instead.
+
+---
+
+## Fallback protocol
+
+### If PDF download fails (403, 404, network error):
+1. Try the ar5iv HTML version: `https://ar5iv.labs.arxiv.org/html/{id}`
+2. If that also fails, try the abstract page: `https://arxiv.org/abs/{id}` to verify the paper exists
+3. If the paper exists but can't be fetched, ask the user to download it manually and provide the path
+
+### If extraction produces garbled text:
+1. Try `pdfplumber` instead of `pymupdf4llm`
+2. If still garbled, fetch ar5iv HTML
+3. If HTML is also bad, try extracting just the text without math preservation
+4. Last resort: ask the user to paste the paper text
+
+### If the paper is very long (>50 pages):
+1. Still extract everything — don't truncate
+2. The section-level files in `sections/` allow reading parts individually
+3. Focus on Method and Appendix sections for implementation details
+
+---
+
+## Quality checklist before proceeding to Stage 2
+
+- [ ] `paper_text.md` exists and is readable
+- [ ] `paper_metadata.json` has title and authors
+- [ ] At least one section file exists in `sections/`
+- [ ] You've checked for appendices
+- [ ] You've checked for algorithm boxes
+- [ ] Equations are present (even if in LaTeX form)
+- [ ] The Method/Model section is identified and readable
+- [ ] You've checked for official code repositories (results in `paper_metadata.json` under `official_code`)
+
+If the Method section is garbled but other sections are fine, attempt to re-extract just that section from the HTML version. The Method section is the most critical — you cannot proceed without a readable version of it.

plugins/codex/plugins/paper2code/skills/paper2code/pipeline/02_contribution_identification.md

@@ -0,0 +1,176 @@
+# Stage 2: Contribution Identification
+
+## Purpose
+Before writing any code, you must understand exactly what this paper contributes and what "implementing this paper" means. This stage forces clarity. 
+
+## Input
+- Parsed paper sections from Stage 1
+- `paper_metadata.json` (includes `official_code` links if any were found)
+
+## Output
+- `.paper2code_work/{ARXIV_ID}/contribution.md` — contains the contribution statement, paper classification, and implementation scope
+
+---
+
+## Reasoning protocol
+
+### Step 1: Read the abstract and introduction 
+
+First read: understand what the paper claims to do.
+Second read: identify the **single sentence** that states the core contribution. This is usually in the abstract ("We propose...," "We introduce...," "In this work, we...") or the last paragraph of the introduction ("Our contributions are...").
+
+Tasks:
+- What is the ONE thing this paper does that didn't exist before?
+- If I had to explain this paper in 10 seconds to an ML engineer, what would I say?
+- What would be missing from ML if this paper didn't exist?
+
+Write this down as a one-sentence contribution statement.
+
+### Step 2: Classify the paper type
+
+Determine which category best fits. This directly affects what you implement:
+
+#### (a) New architecture
+**Examples:** Transformer, ResNet, Vision Transformer, U-Net
+**What "implementing the core contribution" means:** The model architecture — layer definitions, forward pass, how components connect. The architecture IS the contribution.
+**Priority sections:** Model/Architecture section, Figure showing the architecture, any Algorithm boxes
+**What's secondary:** Training procedure (use standard), specific dataset preprocessing
+
+#### (b) New training method or loss function
+**Examples:** DDPM, SimCLR, RLHF, DPO
+**What "implementing the core contribution" means:** The training loop, the loss function, the optimization procedure. The model architecture is usually borrowed from existing work.
+**Priority sections:** Training/Objective section, loss function equations, Algorithm boxes describing the training procedure
+**What's secondary:** Model architecture (often standard), dataset specifics
+
+#### (c) New inference or generation technique
+**Examples:** Beam search variants, nucleus sampling, classifier-free guidance
+**What "implementing the core contribution" means:** The inference algorithm. There may be no training at all.
+**Priority sections:** Inference/Generation section, Algorithm boxes, any sampling or decoding procedures
+**What's secondary:** Training (none or standard), model architecture (usually pre-existing)
+
+#### (d) New dataset or benchmark
+**Examples:** ImageNet, GLUE, MMLU
+**What "implementing the core contribution" means:** The evaluation framework, metrics, data loading pipeline. There may be baseline models but they're not the contribution.
+**Priority sections:** Dataset description, evaluation methodology, metric definitions
+**What's secondary:** Baseline model architectures (reference only)
+
+#### (e) Theoretical analysis with empirical validation
+**Examples:** Lottery Ticket Hypothesis, Scaling Laws
+**What "implementing the core contribution" means:** The experimental setup that validates the theory. Often involves running specific experiments with specific measurements.
+**Priority sections:** Experimental setup, what is measured and how, specific procedures
+**What's secondary:** The theory proofs (important to understand but not to implement)
+
+#### (f) System or engineering paper
+**Examples:** Megatron-LM, FlashAttention, vLLM
+**What "implementing the core contribution" means:** The system design — often involves low-level optimizations, custom kernels, or infrastructure.
+**Priority sections:** System design section, performance benchmarks, implementation details
+**Note:** These papers often cannot be reproduced in a minimal implementation because the contribution IS the engineering. Flag this honestly.
+
+### Step 3: Find the Algorithm box
+
+Search the paper for formal algorithm descriptions (usually labeled "Algorithm 1," "Algorithm 2," etc.). These are gold:
+
+- They are the most precise specification of the procedure
+- They resolve ambiguities in prose descriptions
+- They define variable names and control flow explicitly
+- If the prose says one thing and the Algorithm box says another, **implement the Algorithm box** and flag the discrepancy
+
+If there is no Algorithm box, note this — it means the implementation will rely more on equations and prose, which are more ambiguous.
+
+### Step 4: Check for official code
+
+Check `paper_metadata.json` for the `official_code` field. If official code repositories were found in Stage 1:
+
+1. **Verify the link** — open the URL. Confirm it's the authors' code for this paper, not an unrelated project or empty placeholder.
+2. **Note the framework** — is it PyTorch, JAX, TensorFlow? This may inform your implementation.
+3. **Note the repo structure** — what files exist? This helps you understand the implementation scope.
+4. **Do NOT read the code in detail yet** — that happens in Stage 3 when resolving ambiguities. At this stage, you just need to know it exists and what it covers.
+
+Record the official code status in the contribution statement (Step 5).
+
+### Step 5: Identify what to implement vs. reference
+
+Make three lists:
+
+**IMPLEMENT** (we will write this code):
+- The core contribution (as identified in Step 2)
+- Any component described in enough detail to implement from the paper alone
+
+**REFERENCE** (we will import or note as dependency):
+- Standard components the paper uses but didn't invent (e.g., "standard transformer encoder")
+- Pre-trained models used as backbones
+- Existing techniques referenced with "following [X]" or "similar to [X]"
+
+**OUT OF SCOPE** (we will not touch):
+- Baselines and comparison methods
+- Ablation variants (unless the user specifically requests full mode)
+- Dataset collection or annotation procedures
+- Deployment-specific optimizations
+
+### Step 6: Write the contribution statement
+
+Write a structured contribution statement and save it to `contribution.md`:
+
+```markdown
+# Contribution Analysis
+
+## Paper
+{title} ({year})
+{authors}
+arxiv: {id}
+
+## One-sentence summary
+{what this paper does in one sentence}
+
+## Paper type
+{(a)-(f) classification with justification}
+
+## Core contribution to implement
+{One paragraph describing exactly what we will implement, referencing specific paper sections}
+
+## Algorithm specification
+{Whether an Algorithm box exists, and if so, which one is the primary specification}
+
+## Official code
+{URL if found, or "None found"}
+{Framework used in official code, if applicable}
+{Brief note on what the official repo covers}
+
+## Implementation scope
+### Will implement:
+- {item 1 — with paper section reference}
+- {item 2 — with paper section reference}
+
+### Will reference (not reimplement):
+- {item 1 — what it is, where to get it}
+
+### Out of scope:
+- {item 1 — why}
+
+## Key sections for implementation
+1. {Section X.Y — what it specifies}
+2. {Section X.Y — what it specifies}
+3. {Appendix A — what it specifies}
+```
+
+### Step 7: Self-check
+
+Before proceeding:
+- Could a competent ML engineer read my contribution statement and know exactly what to build? If not, it's too vague.
+- Have I identified the paper type correctly? Read the abstract one more time.
+- Have I checked the appendices for additional specification of the core contribution?
+- Is there anything in my "IMPLEMENT" list that the paper doesn't actually describe in enough detail? Move it to "REFERENCE" or flag it.
+
+---
+
+## Common mistakes at this stage
+
+1. **Trying to implement everything.** A paper about a new loss function doesn't need you to reimplement the backbone architecture. Import it.
+
+2. **Misidentifying the contribution.** A paper might describe a new architecture but the actual contribution is the training recipe. Read what the authors claim, not what takes up the most pages.
+
+3. **Ignoring paper type effects on scope.** The same checklist item (e.g., "training loop") can be in-scope or out-of-scope depending on the paper type. A training method paper MUST include the training loop. An architecture paper doesn't need more than a minimal example.
+
+4. **Skipping the Algorithm box.** If Algorithm 1 exists and you don't read it carefully, you will implement based on noisy prose instead of the precise specification.
+
+5. **Ignoring official code.** Stage 1 automatically searches for code links, but you must verify them in Step 4. If official code exists and you skip it, you'll waste time on ambiguities that the authors already resolved.

plugins/codex/plugins/paper2code/skills/paper2code/pipeline/03_ambiguity_audit.md

@@ -0,0 +1,212 @@
+# Stage 3: The Ambiguity Audit
+
+## Purpose
+This is the most important anti-hallucination stage. Before generating a single line of code, systematically audit every implementation-relevant detail and classify it as SPECIFIED, PARTIALLY_SPECIFIED, or UNSPECIFIED. This stage produces the raw material for `REPRODUCTION_NOTES.md`.
+
+## Input
+- Parsed paper sections from Stage 1
+- Contribution analysis from Stage 2 (includes official code status)
+- Official code repository (if found in Stage 1 — URL in `paper_metadata.json` under `official_code`)
+
+## Output
+- `.paper2code_work/{ARXIV_ID}/ambiguity_audit.md` — complete audit with classification for every item
+
+---
+
+## Critical mindset
+
+**Assume nothing. Verify everything.**
+
+When you read "we use standard settings," ask: standard according to whom? When you read "following prior work," ask: which prior work, and which specific detail from that work? When something seems obvious, ask: is it obvious because it's specified, or because you're making an assumption?
+
+The goal of this stage is to produce a document where a researcher can look at any implementation choice and immediately know whether the paper specified it or whether you chose it.
+
+---
+
+## Reasoning protocol
+
+### The complete checklist
+
+Go through every item below. For each one, search the paper (method section, experiments section, appendix, footnotes — all of them) and classify:
+
+- **SPECIFIED**: The paper states this explicitly. Record the exact quote and section.
+- **PARTIALLY_SPECIFIED**: The paper mentions this but is ambiguous. Record the quote and what's unclear.
+- **UNSPECIFIED**: The paper does not state this at all. Record common choices and which one you'll use as default.
+
+---
+
+### Architecture details
+
+| Item | What to look for | Common hiding spots |
+|------|------------------|---------------------|
+| Layer count / depth | "N layers," "L = 6" | Table 1, first paragraph of experiments |
+| Hidden dimensions | "d_model = 512," "hidden size" | Architecture figure caption, config table in appendix |
+| Number of attention heads | "h = 8 heads" | Same as hidden dims |
+| Head dimension | Often d_model/h, but not always stated | Sometimes only implicit |
+| Feed-forward inner dimension | "d_ff = 2048," "4x expansion" | Appendix, architecture details |
+| Activation functions | "ReLU," "GELU," "SiLU/Swish" | Often unstated — do NOT assume ReLU |
+| Normalization type | LayerNorm, BatchNorm, RMSNorm | Often stated, but placement (pre/post) often unstated |
+| Normalization placement | Pre-norm vs post-norm | CRITICAL for transformers. Check Figure vs text — they often disagree |
+| Normalization epsilon | Almost never stated | Default varies: 1e-5 (PyTorch), 1e-6 (common in papers), 1e-8 (rare) |
+| Initialization scheme | Xavier, Kaiming, custom | Appendix or not stated at all |
+| Dropout rate | "dropout = 0.1" | Appendix, sometimes different rates for different layers |
+| Dropout placement | After attention? After FFN? On embeddings? | Rarely specified precisely |
+| Residual connections | Usually obvious from figures but connection pattern matters | Check if pre-norm or post-norm |
+| Bias terms | Whether linear layers have bias | Almost never stated explicitly |
+| Weight tying | Whether embedding and output projection share weights | Sometimes mentioned casually |
+| Vocabulary size | For NLP models | Dataset description section |
+| Max sequence length | For sequence models | Experiments section, dataset preprocessing |
+
+### Training details
+
+| Item | What to look for | Common hiding spots |
+|------|------------------|---------------------|
+| Optimizer | "Adam," "SGD," "AdamW" | Method or experiments section |
+| Learning rate | The peak/base learning rate | Experiments, appendix |
+| Adam betas | β₁, β₂ values | Appendix if anywhere — papers often say "Adam" without specifying |
+| Adam epsilon | ε value | Almost never stated |
+| Weight decay | Value and where applied | Appendix. Note: often NOT applied to biases and LayerNorm |
+| LR schedule | Warmup type, decay type | Appendix or experiments |
+| Warmup steps | How many steps/epochs of warmup | Appendix — often unstated even when warmup is mentioned |
+| Total training steps/epochs | How long to train | Experiments section |
+| Batch size | Total or per-GPU? | Crucial distinction — often ambiguous |
+| Gradient accumulation | How many micro-batches? | Appendix or unstated |
+| Gradient clipping | Max norm value | Appendix or unstated |
+| Mixed precision | FP16, BF16, or FP32? | Appendix, systems section |
+| Random seed | For reproducibility | Appendix or unstated |
+| Number of GPUs/hardware | Training hardware | Often in experiments or appendix |
+| EMA | Exponential moving average of weights? Decay rate? | Common in diffusion models, often in appendix |
+
+### Data details
+
+| Item | What to look for | Common hiding spots |
+|------|------------------|---------------------|
+| Dataset name and version | Exact dataset identifier | Experiments section |
+| Preprocessing steps | Tokenization, normalization, resizing | Appendix, experiments |
+| Data augmentation | Random crop, flip, color jitter, etc. | Appendix or buried in experiments |
+| Train/val/test split | Standard or custom? | Usually standard but verify |
+| Sequence length / image size | Input dimensions | Experiments, may differ from model max |
+| Tokenizer | Which one specifically (BPE, WordPiece, sentencepiece) | Dataset section, appendix |
+
+### Evaluation details
+
+| Item | What to look for | Common hiding spots |
+|------|------------------|---------------------|
+| Metrics used | Accuracy, BLEU, FID, etc. | Experiments section |
+| Metric computation details | Which BLEU? What tokenizer? Smoothing? | Appendix or reference to external package |
+| Evaluation frequency | Every N steps? Every epoch? | Appendix |
+| Test-time augmentation | Averaging over crops/flips? | Experiments or appendix |
+| Reported numbers | What Table 1 / main results show | Results section — record these for validation |
+
+---
+
+### How to search the paper
+
+For each item in the checklist:
+
+1. **Search the Method section** — the obvious place
+2. **Search the Experiments section** — often contains "Implementation Details" or "Training Details" subsection
+3. **Search the Appendix** — THE most common hiding spot for training details
+4. **Search figure captions** — sometimes specify architectural details not in text
+5. **Search footnotes** — authors sometimes put caveats and clarifications in footnotes
+6. **Search the table captions** — hyperparameter tables often have footnotes with additional info
+7. **Search for the specific term** — Ctrl+F the paper text for "learning rate," "dropout," "optimizer," etc.
+
+---
+
+### How to handle "following X" references
+
+When the paper says "following [X]" or "similar to [X]" or "as in [X]":
+
+1. Note the reference — what is paper X?
+2. If X is a well-known paper (e.g., "following Vaswani et al."), describe what X specifies
+3. If X is an obscure paper, flag this as requiring the reader to look up that paper
+4. Record: `[PARTIALLY_SPECIFIED] Paper says "following [X]" for {component} — X specifies {detail} but reader should verify`
+
+---
+
+### How to resolve ambiguities using official code
+
+If official code was found in Stage 1, use it as a primary resource for resolving `UNSPECIFIED` and `PARTIALLY_SPECIFIED` items. This is the single most effective way to reduce ambiguity.
+
+**Process:**
+1. For each `UNSPECIFIED` item, search the official repo for the relevant implementation detail (e.g., grep for "eps", "dropout", "lr" in config files and model code).
+2. If you find the answer, change the classification:
+   - `UNSPECIFIED` → `SPECIFIED` with tag `[FROM_OFFICIAL_CODE]`
+   - `PARTIALLY_SPECIFIED` → `SPECIFIED` with tag `[FROM_OFFICIAL_CODE]`
+3. Record the exact file and line from the official repo: `github.com/author/repo/blob/main/model.py#L42`
+
+**What official code can resolve:**
+- Activation functions, normalization epsilon, dropout placement
+- Initialization schemes (often the biggest silent assumption)
+- Learning rate schedule details, optimizer parameters
+- Data preprocessing steps
+- Any "standard settings" or "following prior work" references
+
+**What official code cannot resolve:**
+- Whether the code matches the paper's intent (bugs exist in official code too)
+- Whether a detail was chosen for the paper's experiments vs. for engineering convenience
+
+**Important:** Official code is a reference, not ground truth. If the official code contradicts the paper, note both and flag the discrepancy. See the errata section below.
+
+---
+
+### How to handle errata and contradictions
+
+1. **Figure vs text disagreement**: Flag it explicitly. Implement what the text/equations say (they're usually more carefully reviewed than figures) but note the figure shows something different.
+
+2. **Abstract vs method section disagreement**: Method section wins. The abstract is a summary and may oversimplify.
+
+3. **Equation vs prose disagreement**: Equation wins. It's more precise. Flag the discrepancy.
+
+4. **Paper vs official code disagreement**: Note both. The official code may fix bugs or reflect a later understanding, but the paper is what was peer-reviewed. Record: `Paper says X in §Y.Z, but official code does W — we implement {your choice} because {reason}`
+
+---
+
+## Output format
+
+Save to `.paper2code_work/{ARXIV_ID}/ambiguity_audit.md`:
+
+```markdown
+# Ambiguity Audit: {paper_title}
+
+## Architecture
+
+| Component | Status | Paper Quote | Section | Our Choice | Alternatives |
+|-----------|--------|-------------|---------|------------|-------------- |
+| Hidden dim | SPECIFIED | "d_model = 512" | §3.1 | 512 | — |
+| Activation | UNSPECIFIED | — | — | GELU | ReLU, SiLU |
+| Pre/post norm | PARTIALLY_SPECIFIED | "we use layer norm" | §3.2 | Pre-norm | Post-norm (figure suggests post-norm) |
+
+## Training
+{same table format}
+
+## Data
+{same table format}
+
+## Evaluation
+{same table format}
+
+## Contradictions found
+- {description of any inconsistencies}
+
+## References to check
+- {any "following X" references that need verification}
+
+## Official code
+- URL: {if found}
+- Used to resolve: {list of items resolved by reading official code}
+```
+
+---
+
+## Self-check before proceeding
+
+- [ ] Every item in the checklist has been classified (not just the obvious ones)
+- [ ] Every SPECIFIED item has an exact quote and section reference
+- [ ] Every UNSPECIFIED item has a default choice AND alternatives listed
+- [ ] You checked the appendix (not just the method section)
+- [ ] You checked footnotes
+- [ ] You looked for an official code repository
+- [ ] Any "following [X]" references are flagged
+- [ ] Contradictions between text/figures/equations are noted

plugins/codex/plugins/paper2code/skills/paper2code/pipeline/04_code_generation.md

@@ -0,0 +1,314 @@
+# Stage 4: Code Generation
+
+## Purpose
+Generate the implementation repository. Every line of code must be anchored to the paper. Every unspecified choice must be flagged. The output should look like it was written by someone who read the paper three times and implemented it carefully.
+
+## Input
+- Contribution analysis from Stage 2
+- Ambiguity audit from Stage 3
+- Paper sections from Stage 1
+- Official code repository (if found — URL in `paper_metadata.json` under `official_code`)
+
+## Output
+- `{paper_slug}/` directory with all generated files
+
+---
+
+## Pre-generation checklist
+
+Before writing any code, verify:
+- [ ] You have the contribution statement (what exactly are you implementing?)
+- [ ] You have the ambiguity audit (what is specified vs. unspecified?)
+- [ ] You know the paper type (architecture / training method / inference technique / etc.)
+- [ ] You've read `guardrails/scope_enforcement.md` to know what's in/out of scope
+- [ ] You've read the relevant `knowledge/` files for domain-specific gotchas
+- [ ] You've checked `knowledge/paper_to_code_mistakes.md` for relevant pitfalls
+- [ ] You've checked if official code exists (in `paper_metadata.json`) and reviewed any `[FROM_OFFICIAL_CODE]` items in the ambiguity audit
+
+---
+
+## File generation order
+
+Generate files in this order (dependencies flow downward):
+
+1. `configs/base.yaml` — all hyperparameters first, so code can reference them
+2. `src/utils.py` — shared utilities (masking, positional encoding, helpers)
+3. `src/model.py` — architecture
+4. `src/loss.py` — loss functions
+5. `src/data.py` — dataset and dataloader
+6. `src/train.py` — training loop (if in scope)
+7. `src/evaluate.py` — evaluation
+8. `requirements.txt` — dependencies
+9. `REPRODUCTION_NOTES.md` — from the ambiguity audit
+10. `README.md` — project readme
+
+---
+
+## Citation anchoring — mandatory for all code
+
+Every class definition, every non-trivial function, every implementation choice must have a citation comment. Use this format consistently:
+
+```python
+# §3.2 — "We apply layer normalization before each sub-layer"
+# This is the Pre-LN variant (Ba et al., 2016)
+
+# §3.2, Eq. 2 — attention_weights = softmax(QK^T / sqrt(d_k))
+
+# §4.1, Table 2 — d_model = 512, h = 8, d_ff = 2048
+
+# [UNSPECIFIED] Paper does not state epsilon for LayerNorm — using 1e-6
+# Alternatives: 1e-5 (PyTorch default), 1e-8 (some older implementations)
+
+# [PARTIALLY_SPECIFIED] "We use dropout for regularization" (§3.3)
+# Rate of 0.1 stated in §4.1, but placement not specified — applying after attention and FFN
+
+# [ASSUMPTION] Using pre-norm based on "we found pre-norm more stable" (§4.1)
+# Figure 1 shows post-norm, creating ambiguity. We follow the text.
+
+# [FROM_OFFICIAL_CODE] Using learned positional embeddings
+# github.com/author/repo/blob/main/model.py#L42
+```
+
+The `§` symbol is non-negotiable. Use it consistently for section references.
+
+---
+
+## Shape comments — mandatory for all tensor operations
+
+Every tensor operation must have a comment showing the shape transformation:
+
+```python
+x = self.embedding(input_ids)  # (batch, seq_len) -> (batch, seq_len, d_model)
+q = self.W_q(x)  # (batch, seq_len, d_model) -> (batch, seq_len, d_model)
+q = q.view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)  # (batch, n_heads, seq_len, d_k)
+scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)  # (batch, n_heads, seq_len, seq_len)
+```
+
+---
+
+## configs/base.yaml
+
+Every hyperparameter in one place. No magic numbers in code. Format:
+
+```yaml
+# Configuration for {paper_title}
+# Paper: {arxiv_url}
+# All values cited to paper section or flagged as [UNSPECIFIED]
+
+model:
+  d_model: 512          # §3.1 — "d_model = 512"
+  n_heads: 8            # §3.1 — "h = 8"
+  n_layers: 6           # §3.1 — "N = 6"
+  d_ff: 2048            # §3.1 — "d_ff = 2048"
+  dropout: 0.1          # §4.1 — "P_drop = 0.1"
+  activation: "relu"    # [UNSPECIFIED] — Paper does not specify activation in FFN
+  norm_eps: 1.0e-6      # [UNSPECIFIED] — LayerNorm epsilon not stated
+
+training:
+  optimizer: "adam"      # §4.1 — "We use the Adam optimizer"
+  lr: 0.0001            # §4.1 — from the LR schedule formula
+  betas: [0.9, 0.98]    # §4.1 — "β1 = 0.9, β2 = 0.98"
+  eps: 1.0e-9           # §4.1 — "ε = 10^-9"
+  warmup_steps: 4000    # §4.1 — "warmup_steps = 4000"
+  # ... etc
+```
+
+---
+
+## src/model.py
+
+### Structure
+
+```python
+"""
+{Paper title} — Model Architecture
+
+Paper: {arxiv_url}
+Implements: {one-line description of what this file implements}
+
+Section references:
+  §X.Y — {what it describes}
+  §X.Z — {what it describes}
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class ModelConfig:
+    """Configuration for {model name}.
+    
+    All defaults from {paper reference} unless marked [UNSPECIFIED].
+    """
+    # §X.Y — "description from paper"
+    param_name: type = default_value
+
+
+class ComponentA(nn.Module):
+    """§X.Y — Implements {component description from paper}.
+    
+    "{Exact quote from paper describing this component}"
+    """
+    # ...
+
+
+class ComponentB(nn.Module):
+    """§X.Z — Implements {component description from paper}."""
+    # ...
+
+
+class MainModel(nn.Module):
+    """§X — {Paper's name for the model}.
+    
+    Composed of:
+      - ComponentA (§X.Y)
+      - ComponentB (§X.Z)
+    """
+    
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        # Build components
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass following §X.Y description.
+        
+        Args:
+            x: {description} — shape: (batch, ...)
+            
+        Returns:
+            {description} — shape: (batch, ...)
+        """
+        # Implementation with shape comments on every operation
+```
+
+### Rules for model.py
+
+1. **Match paper notation in variable names.** If the paper uses Q, K, V — use `q`, `k`, `v`. If it uses `h` for heads — use `n_heads` (slightly more readable but add comment `# h in paper notation`).
+
+2. **One class per conceptual component.** If the paper describes attention, feed-forward, and the overall block as separate concepts — make them separate classes. Do not put everything in one monolithic class.
+
+3. **Config dataclass, not scattered parameters.** All hyperparameters come from a config dataclass. No `512` literals anywhere.
+
+4. **Forward method mirrors paper flow.** The operations in `forward()` should follow the same order as the paper describes them. A reader should be able to follow along in the paper.
+
+5. **No dead code.** Don't add methods "for convenience" that aren't needed. No `save/load`, no `from_pretrained`, no `num_parameters` — unless the paper contribution involves them.
+
+---
+
+## src/loss.py
+
+### Rules
+
+1. **Implement the exact equation from the paper.** Reference the equation number.
+2. **Note numerical stability considerations.** Log-sum-exp instead of naive softmax, epsilon in denominators, etc.
+3. **If the paper uses a standard loss with modifications, implement only the modifications and wrap the standard loss.**
+4. **Comment where the paper's definition differs from PyTorch's built-in.** This is a common source of bugs.
+
+---
+
+## src/data.py
+
+### Rules
+
+1. **Provide a Dataset class skeleton**, not a complete implementation. The class should show:
+   - What the expected data format is
+   - What preprocessing is needed
+   - What `__getitem__` returns (shapes and types)
+2. **Include clear TODOs** for dataset-specific logic (download path, preprocessing steps)
+3. **Do not hardcode dataset paths.** Use config or arguments.
+4. **Do not download datasets automatically.** Add instructions in docstring.
+
+---
+
+## src/train.py
+
+### Include only if training is in scope (see contribution analysis)
+
+1. **If the contribution IS the training method** (type b paper): This is a primary deliverable. The training loop must precisely follow the paper's algorithm.
+2. **If the contribution is architectural** (type a paper): Provide a minimal training example that instantiates the model and runs one forward/backward pass. Not a full training pipeline.
+3. **Always include** the optimizer setup with all hyperparameters from the config.
+4. **Always include** the learning rate schedule if specified in the paper.
+
+---
+
+## src/evaluate.py
+
+1. **Implement the exact metrics the paper reports in its main results table.**
+2. **Note which metrics package/computation is used** (e.g., sacrebleu vs nltk BLEU — they give different numbers).
+3. **This is metric computation code, not a full eval pipeline.** No data loading, no model loading, just functions that take predictions and targets and return numbers.
+
+---
+
+## src/utils.py
+
+1. **Only include utilities that are shared across multiple files** (e.g., masking functions, positional encoding).
+2. **Do not create a "utils" dumping ground.** If a utility is only used in one file, put it in that file.
+
+---
+
+## requirements.txt
+
+```
+torch>=2.0.0
+pyyaml>=6.0
+numpy>=1.24.0
+# Add paper-specific dependencies below
+```
+
+Pin major versions only. Add comments for why each dependency is needed.
+
+---
+
+## REPRODUCTION_NOTES.md
+
+Transform the ambiguity audit from Stage 3 into a researcher-facing document. Use the template from `scaffolds/reproduction_notes_template.md`. This is a first-class deliverable — not an afterthought. See the scaffold template for the exact structure.
+
+---
+
+## README.md (for the generated project)
+
+Use the template from `scaffolds/readme_template.md`. Should include:
+- Paper title, authors, link
+- What this implementation covers (from contribution statement)
+- Quick-start: how to run the model
+- File structure with one-line descriptions
+- Citation
+
+---
+
+## Code quality final check
+
+Before declaring this stage complete:
+
+- [ ] No magic numbers — every literal is from the config
+- [ ] No missing citations — every class and non-trivial function has a `§` reference
+- [ ] No missing shape comments — every tensor operation has a shape annotation
+- [ ] No silent assumptions — every UNSPECIFIED item has a `[UNSPECIFIED]` comment
+- [ ] No dead code — everything that exists is used
+- [ ] No relative imports — all imports are absolute
+- [ ] Type hints on all function signatures
+- [ ] The model's forward pass mirrors the paper's description order
+- [ ] The config file has a citation comment on every parameter
+- [ ] `REPRODUCTION_NOTES.md` is complete and covers every ambiguity from Stage 3
+
+---
+
+## What NOT to generate
+
+Read this list before writing any code. If you find yourself about to write any of these, stop:
+
+- Do NOT reimplement standard components the paper references but doesn't describe. Import or note the dependency.
+- Do NOT implement distributed training, checkpointing, experiment tracking, or logging infrastructure (unless it IS the paper's contribution).
+- Do NOT implement baseline methods or comparison approaches.
+- Do NOT download or embed datasets.
+- Do NOT write unit tests (this is a reproduction scaffold, not a production library).
+- Do NOT add CLI argument parsing beyond loading a config file.
+- Do NOT add visualization code unless visualization IS the contribution.
+- Do NOT add `if __name__ == "__main__"` blocks with complex logic — keep scripts minimal.

plugins/codex/plugins/paper2code/skills/paper2code/pipeline/05_walkthrough_notebook.md

@@ -0,0 +1,192 @@
+# Stage 5: Walkthrough Notebook
+
+## Purpose
+Generate a Jupyter notebook that connects paper sections to code, making the implementation pedagogically useful. This is not for training — it's for understanding. A researcher should be able to read this notebook and understand exactly how every paper section maps to code.
+
+## Input
+- All generated code from Stage 4
+- Paper sections from Stage 1
+- Contribution analysis from Stage 2
+
+## Output
+- `{paper_slug}/notebooks/walkthrough.ipynb`
+
+---
+
+## Notebook structure
+
+The notebook follows a strict section pattern. For each major component of the implementation:
+
+### Cell pattern (repeat for each component)
+
+#### Cell 1: Paper excerpt (Markdown)
+A quoted block of the relevant paper passage that describes this component. Include the section number and, if applicable, the equation.
+
+```markdown
+## Multi-Head Attention (§3.2)
+
+> "An attention function can be described as mapping a query and a set of key-value
+> pairs to an output, where the query, keys, values, and output are all vectors.
+> The output is computed as a weighted sum of the values, where the weight assigned
+> to each value is computed by a compatibility function of the query with the
+> corresponding key."
+>
+> — §3.2, Attention Is All You Need
+
+**Equation 1:**
+$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$
+```
+
+#### Cell 2: Code (Python)
+The corresponding code from `src/model.py` or wherever this component is defined. Don't just import — inline the relevant class/function so the reader can see it without switching files.
+
+```python
+# From src/model.py — Multi-Head Attention
+
+class MultiHeadAttention(nn.Module):
+    """§3.2 — Multi-Head Attention mechanism."""
+    # ... full implementation
+```
+
+#### Cell 3: Sanity check (Python)
+A runnable check that verifies the component works correctly with toy dimensions:
+
+```python
+# Sanity check: verify output shapes match paper specification
+config = ModelConfig(d_model=64, n_heads=4)  # Toy dimensions for CPU
+attn = MultiHeadAttention(config)
+
+batch, seq_len = 2, 10
+x = torch.randn(batch, seq_len, config.d_model)
+
+output = attn(x)
+assert output.shape == (batch, seq_len, config.d_model), \
+    f"Expected (2, 10, 64), got {output.shape}"
+print(f"✓ MultiHeadAttention: input {x.shape} → output {output.shape}")
+```
+
+---
+
+## Required sections
+
+### 1. Setup and imports
+```python
+import torch
+import torch.nn as nn
+import math
+# ... minimal imports
+
+# Use tiny dimensions so everything runs on CPU quickly
+BATCH_SIZE = 2
+SEQ_LEN = 16
+D_MODEL = 64  # Paper uses 512, we use 64 for walkthrough
+# ... etc
+```
+
+Explain why dimensions are small: "We use reduced dimensions so the notebook runs instantly on CPU. The architecture is identical — only the numbers change."
+
+### 2. Configuration
+Show the config dataclass with paper citations on each field. Instantiate with toy dimensions.
+
+### 3. One section per architectural component
+Follow the paper's presentation order. If the paper describes components A, B, C used to build the model, the notebook should have sections for A, B, C in that order, then the composed model.
+
+For each: paper excerpt → code → sanity check.
+
+### 4. Full model assembly
+Show how the components compose into the full model. Run a full forward pass with random input. Verify output shape.
+
+```python
+# Full model forward pass
+model = TheModel(config)
+x = torch.randn(BATCH_SIZE, SEQ_LEN)  # or appropriate input
+output = model(x)
+print(f"✓ Full model: input {x.shape} → output {output.shape}")
+print(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")
+```
+
+### 5. Loss function (if in scope)
+Paper excerpt for the loss function → code → compute loss on random predictions.
+
+### 6. Training step (if in scope)
+Show one forward + backward pass. Verify gradients flow:
+
+```python
+# One training step
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
+loss = loss_fn(output, target)
+loss.backward()
+optimizer.step()
+
+# Verify gradients exist
+for name, param in model.named_parameters():
+    if param.grad is None:
+        print(f"⚠ No gradient for: {name}")
+    else:
+        print(f"✓ {name}: grad norm = {param.grad.norm():.4f}")
+```
+
+### 7. Common pitfalls (Markdown)
+A cell listing what typically goes wrong when implementing this paper:
+
+```markdown
+## Common Pitfalls
+
+1. **Pre-norm vs post-norm**: The paper's figure shows post-norm but experiments
+   use pre-norm. This significantly affects training stability. [§3.2]
+
+2. **Attention scale factor**: Must divide by √d_k, not √d_model. If using
+   multi-head attention, d_k = d_model / n_heads. [§3.2, Eq. 1]
+
+3. **Positional encoding**: The sinusoidal encoding uses specific frequencies.
+   Off-by-one errors in the position index are common. [§3.5]
+
+4. **Label smoothing**: The paper uses ε=0.1 label smoothing. PyTorch's
+   CrossEntropyLoss does not support this directly — you need a custom
+   implementation. [§5.4]
+
+5. **Learning rate schedule**: The warmup schedule has a specific formula.
+   Using a generic linear warmup gives different results. [§5.3]
+```
+
+---
+
+## Notebook generation rules
+
+1. **Everything must run on CPU.** No GPU required. Use tiny dimensions.
+2. **No external data needed.** All inputs are random tensors.
+3. **No training to convergence.** Just verify shapes, gradients, and basic numerics.
+4. **Each cell should be self-contained enough to understand.** Imports can be at the top, but each section should be readable on its own.
+5. **Markdown cells should quote the paper directly.** Don't paraphrase — use exact quotes with section references.
+6. **Sanity checks should be assertions, not just prints.** Assert the expected shapes so failures are loud.
+7. **Use the paper's variable names in code.** If the paper uses Q, K, V — the notebook code should too.
+
+---
+
+## Mode-specific behavior
+
+### minimal mode
+Standard notebook as described above.
+
+### educational mode
+Add additional cells:
+- Before each component: a "Background" markdown cell explaining the ML concept (what is attention? what is a residual connection?) for readers who are learning
+- After each sanity check: a "What would happen if..." cell exploring edge cases (what if you remove the scale factor? what if you use post-norm instead of pre-norm?)
+- A "Further Reading" section at the end pointing to tutorials and explanations
+
+### full mode
+Same as minimal but includes:
+- Full data loading example (with mock data)
+- Complete training loop (3-5 steps, not to convergence)
+- Evaluation metric computation on random predictions
+
+---
+
+## Self-check
+
+- [ ] Notebook runs end-to-end without errors (all cells execute)
+- [ ] All assertion-based sanity checks pass
+- [ ] Every paper section relevant to the implementation is quoted
+- [ ] Toy dimensions are used throughout (runs on CPU in seconds)
+- [ ] Common pitfalls section exists and is genuinely useful
+- [ ] No external data or GPU required

plugins/codex/plugins/paper2code/skills/paper2code/scaffolds/config_template.yaml

@@ -0,0 +1,62 @@
+# Configuration for {{PAPER_TITLE}}
+# Paper: https://arxiv.org/abs/{{ARXIV_ID}}
+#
+# CONVENTION:
+#   - Lines with §X.Y cite the paper section specifying this value
+#   - Lines with [UNSPECIFIED] mark choices not stated in the paper
+#   - Lines with [FROM_OFFICIAL_CODE] are resolved from the authors' implementation
+#
+# All hyperparameters in one place. No magic numbers in code.
+
+# ---------------------------------------------------------------------------
+# Model architecture
+# ---------------------------------------------------------------------------
+model:
+  # §X.Y — "quote from paper specifying this value"
+  # param_name: value
+
+  # [UNSPECIFIED] — Paper does not state this. Common choices: A, B, C
+  # param_name: value
+
+# ---------------------------------------------------------------------------
+# Training
+# ---------------------------------------------------------------------------
+training:
+  # §X.Y — "quote from paper"
+  # optimizer: adam
+  # lr: 0.0001
+  # betas: [0.9, 0.999]       # [UNSPECIFIED] — using PyTorch Adam defaults
+  # eps: 1.0e-8               # [UNSPECIFIED] — using PyTorch Adam defaults
+  # weight_decay: 0.0         # [UNSPECIFIED] — no weight decay mentioned
+
+  # Learning rate schedule
+  # schedule: cosine           # §X.Y — "we use cosine annealing"
+  # warmup_steps: 1000         # §X.Y — "linear warmup for 1000 steps"
+  # total_steps: 100000        # §X.Y — "we train for 100k steps"
+
+  # Regularization
+  # dropout: 0.1               # §X.Y — "dropout rate of 0.1"
+  # gradient_clip: 1.0         # [UNSPECIFIED] — common default
+
+  # Batch
+  # batch_size: 256            # §X.Y — "batch size of 256"
+  # batch_size_note: "total"   # §X.Y — "total batch size across 8 GPUs" or [UNSPECIFIED]
+
+  # Hardware (for reference, not used in code)
+  # gpus: 8                    # §X.Y — "we use 8 V100 GPUs"
+  # precision: fp32            # [UNSPECIFIED] — mixed precision not mentioned
+
+# ---------------------------------------------------------------------------
+# Data
+# ---------------------------------------------------------------------------
+data:
+  # dataset: "{{DATASET_NAME}}"    # §X.Y — "we evaluate on {{DATASET}}"
+  # seq_length: 512                # §X.Y — "maximum sequence length of 512"
+  # vocab_size: 30000              # §X.Y or [UNSPECIFIED]
+
+# ---------------------------------------------------------------------------
+# Evaluation
+# ---------------------------------------------------------------------------
+eval:
+  # metric: "accuracy"             # §X.Y — "we report top-1 accuracy"
+  # eval_every: 1000               # [UNSPECIFIED] — evaluation frequency not stated

plugins/codex/plugins/paper2code/skills/paper2code/scaffolds/data_template.py

@@ -0,0 +1,113 @@
+"""
+{{PAPER_TITLE}} — Dataset and Data Loading
+
+Paper: https://arxiv.org/abs/{{ARXIV_ID}}
+Implements: Data loading for {{DATASET_NAME}}
+
+Section references:
+  {{§SECTION}} — {{data description}}
+
+NOTE: This file provides the Dataset class skeleton. You must:
+  1. Download the dataset from {{DATASET_URL}}
+  2. Set the data_dir in configs/base.yaml
+  3. Implement any dataset-specific preprocessing (marked with TODO)
+"""
+
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+
+class {{DATASET_CLASS}}(Dataset):
+    """§{{SECTION}} — Dataset for {{PAPER_TITLE}}.
+    
+    "{{Quote from paper about the dataset used}}"
+    
+    Expected data format:
+        {{describe the expected file structure / data format}}
+    
+    How to obtain:
+        {{instructions for downloading the dataset}}
+    
+    Preprocessing:
+        {{describe preprocessing steps from the paper}}
+    """
+    
+    def __init__(
+        self,
+        data_dir: str,
+        split: str = "train",
+        # Add other params from config
+    ):
+        """
+        Args:
+            data_dir: path to the dataset root directory
+            split: one of "train", "val", "test"
+        """
+        self.data_dir = Path(data_dir)
+        self.split = split
+        
+        # TODO: Load file list / metadata
+        # self.samples = self._load_samples()
+    
+    def _load_samples(self):
+        """Load sample paths/metadata for the given split.
+        
+        TODO: Implement based on the dataset structure.
+        """
+        raise NotImplementedError(
+            f"Dataset loading not implemented. "
+            f"Download the dataset and implement _load_samples() for your data format."
+        )
+    
+    def _preprocess(self, sample):
+        """Apply preprocessing as described in §{{SECTION}}.
+        
+        TODO: Implement the paper's preprocessing pipeline:
+        {{list preprocessing steps from the paper}}
+        """
+        raise NotImplementedError("Implement preprocessing per §{{SECTION}}")
+    
+    def __len__(self) -> int:
+        return len(self.samples)
+    
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """Load and preprocess a single sample.
+        
+        Returns:
+            dict with keys:
+                {{key_1}}: {{description}} — shape: {{shape}}
+                {{key_2}}: {{description}} — shape: {{shape}}
+        """
+        # TODO: Implement sample loading
+        # sample = self.samples[idx]
+        # processed = self._preprocess(sample)
+        # return processed
+        raise NotImplementedError("Implement __getitem__ for your data format")
+
+
+def build_dataloader(
+    config: dict,
+    split: str = "train",
+) -> DataLoader:
+    """Build a DataLoader from config.
+    
+    Args:
+        config: data config dict from base.yaml
+        split: "train", "val", or "test"
+    """
+    dataset = {{DATASET_CLASS}}(
+        data_dir=config["data_dir"],
+        split=split,
+    )
+    
+    return DataLoader(
+        dataset,
+        batch_size=config.get("batch_size", 32),
+        shuffle=(split == "train"),
+        num_workers=config.get("num_workers", 4),
+        pin_memory=True,
+        drop_last=(split == "train"),
+    )

plugins/codex/plugins/paper2code/skills/paper2code/scaffolds/evaluate_template.py

@@ -0,0 +1,62 @@
+"""
+{{PAPER_TITLE}} — Evaluation Metrics
+
+Paper: https://arxiv.org/abs/{{ARXIV_ID}}
+Implements: Metric computation for {{METRICS_LISTED}}
+
+Section references:
+  {{§SECTION}} — {{evaluation methodology description}}
+
+NOTE: This file provides metric computation functions only.
+      It does NOT handle data loading or model inference.
+      Pass predictions and targets to compute metrics.
+"""
+
+from typing import Dict, List, Optional
+
+import torch
+import numpy as np
+
+
+def compute_{{METRIC_1}}(
+    predictions: torch.Tensor,
+    targets: torch.Tensor,
+) -> float:
+    """§{{SECTION}} — Compute {{metric name}}.
+    
+    "{{Quote from paper about evaluation metric}}"
+    
+    Args:
+        predictions: model output — shape: (batch, {{dims}})
+        targets: ground truth — shape: (batch, {{dims}})
+        
+    Returns:
+        {{metric_name}} score as a float
+        
+    NOTE: {{any caveats about metric implementation, e.g.,
+           "Different BLEU implementations give different numbers.
+            We use sacrebleu for reproducibility."}}
+    """
+    # §{{SECTION}} — metric computation
+    pass  # REPLACE with actual metric
+
+
+def compute_all_metrics(
+    predictions: torch.Tensor,
+    targets: torch.Tensor,
+) -> Dict[str, float]:
+    """Compute all evaluation metrics reported in the paper.
+    
+    §{{SECTION}} — "We report {{metric_1}}, {{metric_2}}, and {{metric_3}}"
+    
+    Args:
+        predictions: model output — shape: (batch, {{dims}})
+        targets: ground truth — shape: (batch, {{dims}})
+    
+    Returns:
+        Dict mapping metric name to value
+    """
+    return {
+        "{{metric_1}}": compute_{{METRIC_1}}(predictions, targets),
+        # Add more metrics as needed
+    }

plugins/codex/plugins/paper2code/skills/paper2code/scaffolds/loss_template.py

@@ -0,0 +1,53 @@
+"""
+{{PAPER_TITLE}} — Loss Functions
+
+Paper: https://arxiv.org/abs/{{ARXIV_ID}}
+Implements: {{LOSS_DESCRIPTION}}
+
+Section references:
+  {{§SECTION}}, Eq. {{N}} — {{DESCRIPTION}}
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class {{LOSS_CLASS}}(nn.Module):
+    """§{{SECTION}}, Eq. {{N}} — {{Loss name from paper}}.
+    
+    "{{Exact quote from paper describing the loss function}}"
+    
+    Mathematical formulation:
+        {{L = formula as written in paper}}
+    
+    Args:
+        {{param}}: {{description}} — §{{SECTION}} specifies {{value}}
+    """
+    
+    def __init__(self, {{params}}):
+        super().__init__()
+        # §{{SECTION}} — loss hyperparameters
+        pass  # REPLACE with actual parameters
+    
+    def forward(
+        self,
+        predictions: torch.Tensor,
+        targets: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute loss.
+        
+        Args:
+            predictions: model output — shape: (batch, {{dims}})
+            targets: ground truth — shape: (batch, {{dims}})
+            
+        Returns:
+            Scalar loss value
+        """
+        # §{{SECTION}}, Eq. {{N}} — step-by-step following the equation
+        # Each intermediate computation gets a shape comment
+        
+        # NOTE on numerical stability:
+        # {{describe any stability considerations}}
+        
+        pass  # REPLACE with actual loss computation

plugins/codex/plugins/paper2code/skills/paper2code/scaffolds/model_template.py

@@ -0,0 +1,139 @@
+"""
+{{PAPER_TITLE}} — Model Architecture
+
+Paper: https://arxiv.org/abs/{{ARXIV_ID}}
+Authors: {{AUTHORS}}
+Year: {{YEAR}}
+
+Implements: {{ONE_LINE_DESCRIPTION}}
+
+Section references:
+  {{§SECTION_1}} — {{DESCRIPTION_1}}
+  {{§SECTION_2}} — {{DESCRIPTION_2}}
+  {{§SECTION_3}} — {{DESCRIPTION_3}}
+
+Usage:
+    from src.model import {{MODEL_CLASS}}, ModelConfig
+    
+    config = ModelConfig()
+    model = {{MODEL_CLASS}}(config)
+    output = model(input_tensor)
+"""
+
+import math
+from dataclasses import dataclass, field
+from typing import Optional, Tuple, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+@dataclass
+class ModelConfig:
+    """All model hyperparameters.
+    
+    Values from {{PAPER_TITLE}} unless marked [UNSPECIFIED].
+    Matches configs/base.yaml — change values there, not here.
+    """
+    # Architecture — §{{ARCH_SECTION}}
+    # {{PARAM_1_NAME}}: {{TYPE}} = {{VALUE}}  # §X.Y — "quote from paper"
+    # {{PARAM_2_NAME}}: {{TYPE}} = {{VALUE}}  # [UNSPECIFIED] — our choice, alternatives: ...
+
+    pass  # REPLACE with actual config fields
+
+
+# ---------------------------------------------------------------------------
+# Sub-modules
+# ---------------------------------------------------------------------------
+
+class {{COMPONENT_A}}(nn.Module):
+    """§{{SECTION}} — {{Description of component from paper}}.
+    
+    "{{Exact quote from paper describing this component}}"
+    """
+    
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        # §{{SECTION}} — build layers as described
+        pass  # REPLACE with actual layers
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: {{description}} — shape: (batch, {{dims}})
+            
+        Returns:
+            {{description}} — shape: (batch, {{dims}})
+        """
+        # §{{SECTION}} — forward pass
+        # Every tensor operation gets a shape comment:
+        # x = self.linear(x)  # (batch, seq_len, d_model) -> (batch, seq_len, d_ff)
+        pass  # REPLACE with actual forward pass
+
+
+class {{COMPONENT_B}}(nn.Module):
+    """§{{SECTION}} — {{Description of component from paper}}."""
+    
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        pass  # REPLACE
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pass  # REPLACE
+
+
+# ---------------------------------------------------------------------------
+# Main Model
+# ---------------------------------------------------------------------------
+
+class {{MODEL_CLASS}}(nn.Module):
+    """§{{SECTION}} — {{Paper's name for the full model}}.
+    
+    Composed of:
+      - {{COMPONENT_A}} (§{{SECTION_A}})
+      - {{COMPONENT_B}} (§{{SECTION_B}})
+    
+    "{{Quote from paper describing the overall model}}"
+    """
+    
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        
+        # Build model components
+        # REPLACE with actual component instantiation
+    
+    def forward(
+        self,
+        x: torch.Tensor,
+        # Add other inputs as needed (mask, labels, etc.)
+    ) -> torch.Tensor:
+        """Forward pass following §{{SECTION}} description.
+        
+        Args:
+            x: {{description}} — shape: (batch, {{input_dims}})
+            
+        Returns:
+            {{description}} — shape: (batch, {{output_dims}})
+        """
+        # §{{SECTION}} — step-by-step forward pass
+        # Mirror the paper's description order
+        # Shape comments on every operation
+        pass  # REPLACE
+    
+    def __repr__(self) -> str:
+        """Print architecture summary."""
+        total_params = sum(p.numel() for p in self.parameters())
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"  config={self.config},\n"
+            f"  total_params={total_params:,},\n"
+            f"  trainable_params={trainable_params:,}\n"
+            f")"
+        )

plugins/codex/plugins/paper2code/skills/paper2code/scaffolds/readme_template.md

@@ -0,0 +1,62 @@
+# {{PAPER_TITLE}}
+
+Implementation of [{{PAPER_TITLE}}](https://arxiv.org/abs/{{ARXIV_ID}}) ({{AUTHORS}}, {{YEAR}}).
+
+## What this implements
+
+{{ONE_PARAGRAPH_CONTRIBUTION_STATEMENT}}
+
+## Quick start
+
+```bash
+pip install -r requirements.txt
+```
+
+```python
+from src.model import {{MODEL_CLASS}}, ModelConfig
+
+config = ModelConfig()
+model = {{MODEL_CLASS}}(config)
+
+# Example forward pass
+import torch
+x = torch.randn({{EXAMPLE_INPUT_SHAPE}})
+output = model(x)
+print(output.shape)  # {{EXPECTED_OUTPUT_SHAPE}}
+```
+
+## File structure
+
+```
+{{PAPER_SLUG}}/
+├── README.md                 # This file
+├── REPRODUCTION_NOTES.md     # Ambiguity audit — what's specified vs. assumed
+├── requirements.txt          # Dependencies
+├── src/
+│   ├── model.py              # {{MODEL_DESCRIPTION}}
+│   ├── loss.py               # {{LOSS_DESCRIPTION}}
+│   ├── data.py               # {{DATA_DESCRIPTION}}
+│   ├── train.py              # {{TRAIN_DESCRIPTION}}
+│   ├── evaluate.py           # {{EVAL_DESCRIPTION}}
+│   └── utils.py              # {{UTILS_DESCRIPTION}}
+├── configs/
+│   └── base.yaml             # All hyperparameters — each cited or flagged
+└── notebooks/
+    └── walkthrough.ipynb     # Paper sections → code → sanity checks
+```
+
+## Important: Read REPRODUCTION_NOTES.md
+
+This implementation flags every choice that the paper does not specify.
+Before using this code for research, read [REPRODUCTION_NOTES.md](REPRODUCTION_NOTES.md)
+to understand which implementation details are from the paper and which are our choices.
+
+## Citation
+
+```bibtex
+{{BIBTEX_ENTRY}}
+```
+
+---
+
+*Generated by [paper2code](https://github.com/PrathamLearnsToCode/paper2code) — citation-anchored paper implementation.*

plugins/codex/plugins/paper2code/skills/paper2code/scaffolds/reproduction_notes_template.md

@@ -0,0 +1,109 @@
+# Reproduction Notes: {{PAPER_TITLE}}
+
+> This document records every implementation choice, whether it was specified by the paper,
+> and what alternatives exist. If you're reproducing this paper, **read this first.**
+
+---
+
+## Paper
+
+- **Title:** {{PAPER_TITLE}}
+- **Authors:** {{AUTHORS}}
+- **Year:** {{YEAR}}
+- **ArXiv:** https://arxiv.org/abs/{{ARXIV_ID}}
+- **Official code:** {{OFFICIAL_CODE_URL or "None found"}}
+
+---
+
+## What this implements
+
+{{ONE_PARAGRAPH_CONTRIBUTION_STATEMENT — from Stage 2}}
+
+---
+
+## Verified against
+
+{{Describe what the implementation has been checked against:}}
+- [ ] Paper equations (§X.Y, Eq. N)
+- [ ] Paper Algorithm box (Algorithm 1)
+- [ ] Official code (if available — link specific files/lines)
+- [ ] Well-known reimplementation (name, link)
+- [ ] None — implementation based solely on paper text
+
+---
+
+## Unspecified choices
+
+Every implementation choice that the paper does not explicitly specify.
+
+| Component | Our Choice | Alternatives | Paper Quote (if partial) | Section |
+|-----------|-----------|--------------|--------------------------|---------|
+| {{component}} | {{what we used}} | {{other options}} | {{quote or "—"}} | {{§X.Y or "—"}} |
+
+<!-- Example rows: -->
+<!-- | LayerNorm epsilon | 1e-6 | 1e-5 (PyTorch default), 1e-8 | — | — | -->
+<!-- | Activation in FFN | GELU | ReLU, SiLU | — | — | -->
+<!-- | Weight initialization | PyTorch defaults | Xavier, normal(0, 0.02) | — | — | -->
+<!-- | Dropout placement | After attn + FFN | Various | "We use dropout" | §3.3 | -->
+
+---
+
+## Known deviations
+
+Anything this implementation intentionally does differently from the paper, and why.
+
+| Deviation | Paper says | We do | Reason |
+|-----------|-----------|-------|--------|
+| {{what}} | {{paper specification}} | {{our implementation}} | {{why we deviate}} |
+
+<!-- Example: -->
+<!-- | Figure vs text | Figure 1 shows post-norm | We use pre-norm | Text §3.2 says "pre-norm", Figure is likely outdated | -->
+
+---
+
+## Expected results
+
+What metrics should you see if the implementation is correct? From the paper's main results.
+
+| Metric | Paper's number | Dataset | Conditions |
+|--------|---------------|---------|------------|
+| {{metric}} | {{value}} | {{dataset}} | {{e.g., "Table 1, best config"}} |
+
+**Note:** Exact reproduction of paper numbers requires matching ALL unspecified choices above,
+plus having the exact training data, hardware, and random seeds. Small deviations (0.1-0.5%)
+are normal even with correct implementations.
+
+---
+
+## Debugging tips
+
+Common failure modes when reproducing this type of paper:
+
+1. **{{Failure mode 1}}**: {{description and how to diagnose}}
+2. **{{Failure mode 2}}**: {{description and how to diagnose}}
+3. **{{Failure mode 3}}**: {{description and how to diagnose}}
+
+<!-- Example: -->
+<!-- 1. **Loss not decreasing**: Check pre-norm vs post-norm. Post-norm without warmup often diverges. -->
+<!-- 2. **NaN in attention**: Masking with 0 instead of -inf can cause NaN after softmax with all-zero rows. -->
+<!-- 3. **Lower accuracy than paper**: Check if label smoothing is implemented — it adds ~0.5% for transformers. -->
+
+---
+
+## Scope decisions
+
+### Implemented
+- {{Component}} — {{reason: core contribution / necessary for X}}
+
+### Intentionally excluded
+- {{Component}} — {{reason: baseline method / standard component / not in paper scope}}
+
+### Needed for full reproduction (not included)
+- {{Component}} — {{what it is, where to get it}}
+
+---
+
+## References
+
+Papers referenced by this implementation:
+- {{Citation}} — {{what was taken from this reference}}

plugins/codex/plugins/paper2code/skills/paper2code/scaffolds/train_template.py

@@ -0,0 +1,183 @@
+"""
+{{PAPER_TITLE}} — Training Loop
+
+Paper: https://arxiv.org/abs/{{ARXIV_ID}}
+Implements: {{TRAINING_DESCRIPTION}}
+
+Section references:
+  {{§SECTION}} — {{DESCRIPTION}}
+"""
+
+import math
+import yaml
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from src.model import {{MODEL_CLASS}}, ModelConfig
+from src.loss import {{LOSS_FN}}
+from src.data import {{DATASET_CLASS}}
+
+
+def load_config(config_path: str = "configs/base.yaml") -> dict:
+    """Load training configuration."""
+    with open(config_path) as f:
+        return yaml.safe_load(f)
+
+
+def build_optimizer(
+    model: nn.Module,
+    config: dict,
+) -> torch.optim.Optimizer:
+    """Build optimizer with parameter groups.
+    
+    §{{SECTION}} — "{{quote about optimizer}}"
+    """
+    # [ASSUMPTION] Not applying weight decay to biases and normalization layers
+    # Paper does not specify which parameters get weight decay
+    no_decay = ["bias", "LayerNorm.weight", "LayerNorm.bias"]
+    param_groups = [
+        {
+            "params": [
+                p for n, p in model.named_parameters()
+                if not any(nd in n for nd in no_decay) and p.requires_grad
+            ],
+            "weight_decay": config["training"]["weight_decay"],
+        },
+        {
+            "params": [
+                p for n, p in model.named_parameters()
+                if any(nd in n for nd in no_decay) and p.requires_grad
+            ],
+            "weight_decay": 0.0,
+        },
+    ]
+
+    optimizer_name = config["training"]["optimizer"]
+    lr = config["training"]["lr"]
+
+    if optimizer_name == "adam":
+        # §{{SECTION}} — optimizer specification
+        return torch.optim.Adam(
+            param_groups,
+            lr=lr,
+            betas=tuple(config["training"].get("betas", [0.9, 0.999])),
+            eps=config["training"].get("eps", 1e-8),
+        )
+    elif optimizer_name == "adamw":
+        return torch.optim.AdamW(
+            param_groups,
+            lr=lr,
+            betas=tuple(config["training"].get("betas", [0.9, 0.999])),
+            eps=config["training"].get("eps", 1e-8),
+        )
+    elif optimizer_name == "sgd":
+        return torch.optim.SGD(
+            param_groups,
+            lr=lr,
+            momentum=config["training"].get("momentum", 0.9),
+        )
+    else:
+        raise ValueError(f"Unknown optimizer: {optimizer_name}")
+
+
+def build_scheduler(
+    optimizer: torch.optim.Optimizer,
+    config: dict,
+) -> Optional[torch.optim.lr_scheduler._LRScheduler]:
+    """Build learning rate scheduler.
+    
+    §{{SECTION}} — "{{quote about LR schedule}}"
+    """
+    schedule = config["training"].get("schedule", "constant")
+    warmup_steps = config["training"].get("warmup_steps", 0)
+    total_steps = config["training"]["total_steps"]
+
+    if schedule == "constant":
+        return None
+    elif schedule == "cosine":
+        def lr_lambda(step):
+            if step < warmup_steps:
+                return step / max(warmup_steps, 1)
+            progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
+            return 0.5 * (1.0 + math.cos(math.pi * progress))
+        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+    elif schedule == "linear":
+        def lr_lambda(step):
+            if step < warmup_steps:
+                return step / max(warmup_steps, 1)
+            return max(0.0, (total_steps - step) / max(total_steps - warmup_steps, 1))
+        return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+    else:
+        raise ValueError(f"Unknown schedule: {schedule}")
+
+
+def train(config_path: str = "configs/base.yaml"):
+    """Main training loop.
+    
+    {{Describe what this training loop does and which paper sections it follows.}}
+    """
+    config = load_config(config_path)
+
+    # Device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Build model
+    model_config = ModelConfig(
+        # REPLACE: populate from config dict
+    )
+    model = {{MODEL_CLASS}}(model_config).to(device)
+
+    # Build optimizer and scheduler
+    optimizer = build_optimizer(model, config)
+    scheduler = build_scheduler(optimizer, config)
+
+    # Build loss
+    loss_fn = {{LOSS_FN}}()  # REPLACE with actual loss construction
+
+    # Build data — see src/data.py for dataset setup instructions
+    # dataset = {{DATASET_CLASS}}(config["data"])
+    # dataloader = DataLoader(dataset, batch_size=config["training"]["batch_size"],
+    #                         shuffle=True, num_workers=4, pin_memory=True)
+
+    # Training loop
+    # §{{SECTION}} — training procedure
+    gradient_clip = config["training"].get("gradient_clip", None)
+    total_steps = config["training"]["total_steps"]
+
+    model.train()
+    step = 0
+
+    print(f"Starting training for {total_steps} steps")
+    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+    # REPLACE: implement the actual training loop
+    # for epoch in range(max_epochs):
+    #     for batch in dataloader:
+    #         batch = {k: v.to(device) for k, v in batch.items()}
+    #
+    #         optimizer.zero_grad()
+    #         output = model(batch["input"])
+    #         loss = loss_fn(output, batch["target"])
+    #         loss.backward()
+    #
+    #         if gradient_clip is not None:
+    #             torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clip)
+    #
+    #         optimizer.step()
+    #         if scheduler is not None:
+    #             scheduler.step()
+    #
+    #         step += 1
+    #         if step % 100 == 0:
+    #             print(f"Step {step}/{total_steps} | Loss: {loss.item():.4f}")
+    #
+    #         if step >= total_steps:
+    #             break
+
+
+if __name__ == "__main__":
+    train()

plugins/codex/plugins/paper2code/skills/paper2code/scripts/extract_structure.py

@@ -0,0 +1,335 @@
+#!/usr/bin/env python3
+"""
+Extract structure from a parsed paper text.
+
+Usage:
+    python extract_structure.py <paper_text.md> <output_dir>
+
+Outputs:
+    {output_dir}/sections/         — individual section files
+    {output_dir}/algorithms/       — extracted algorithm boxes
+    {output_dir}/equations/        — extracted numbered equations
+    {output_dir}/tables/           — extracted tables
+    {output_dir}/footnotes.md      — all footnotes collected
+"""
+
+import re
+import sys
+from pathlib import Path
+
+
+def identify_sections(text: str) -> list[dict]:
+    """Identify section boundaries using heading patterns.
+
+    Detects:
+      - Markdown headings (# , ## , ### )
+      - Numbered headings (1. Introduction, 2.1 Related Work)
+      - ALL CAPS headings (INTRODUCTION, RELATED WORK)
+    """
+    lines = text.split("\n")
+    sections = []
+    current_section = None
+    current_lines = []
+
+    # Patterns for section headings
+    md_heading = re.compile(r"^(#{1,4})\s+(.+)$")
+    numbered_heading = re.compile(
+        r"^(\d+(?:\.\d+)*)\s+([A-Z][A-Za-z\s:,\-]+)$"
+    )
+    allcaps_heading = re.compile(r"^([A-Z][A-Z\s]{4,})$")
+
+    def save_current():
+        if current_section and current_lines:
+            sections.append({
+                "title": current_section,
+                "content": "\n".join(current_lines).strip(),
+            })
+
+    for line in lines:
+        heading = None
+
+        # Check markdown heading
+        m = md_heading.match(line)
+        if m:
+            heading = m.group(2).strip()
+
+        # Check numbered heading
+        if not heading:
+            m = numbered_heading.match(line.strip())
+            if m:
+                heading = f"{m.group(1)} {m.group(2).strip()}"
+
+        # Check ALL CAPS heading (only for longer titles to avoid false positives)
+        if not heading:
+            m = allcaps_heading.match(line.strip())
+            if m and len(m.group(1).strip()) > 5:
+                heading = m.group(1).strip().title()
+
+        if heading:
+            save_current()
+            current_section = heading
+            current_lines = []
+        else:
+            current_lines.append(line)
+
+    save_current()
+    return sections
+
+
+def extract_algorithms(text: str) -> list[dict]:
+    """Extract algorithm boxes from the paper.
+
+    Looks for patterns like:
+      Algorithm 1: Name
+      ...algorithm body...
+      (ends at next section heading or next Algorithm block)
+    """
+    algorithms = []
+
+    # Pattern: "Algorithm N" possibly followed by colon and name
+    pattern = re.compile(
+        r"(Algorithm\s+\d+[:\.]?\s*[^\n]*)\n(.*?)(?=Algorithm\s+\d+[:\.]|^#{1,4}\s|\Z)",
+        re.DOTALL | re.MULTILINE,
+    )
+
+    for match in pattern.finditer(text):
+        title = match.group(1).strip()
+        body = match.group(2).strip()
+        if body:
+            algorithms.append({
+                "title": title,
+                "content": body,
+            })
+
+    return algorithms
+
+
+def extract_equations(text: str) -> list[dict]:
+    """Extract numbered equations.
+
+    Looks for:
+      - LaTeX equation environments: \\begin{equation}...\\end{equation}
+      - Display math with numbering: $$ ... $$ (N)
+      - Inline equation references: (1), (2), Eq. 1, Equation 1
+      - Markdown math blocks
+    """
+    equations = []
+
+    # LaTeX equation environments
+    latex_eq = re.compile(
+        r"\\begin\{(?:equation|align|gather)\*?\}(.*?)\\end\{(?:equation|align|gather)\*?\}",
+        re.DOTALL,
+    )
+    for i, match in enumerate(latex_eq.finditer(text)):
+        equations.append({
+            "number": i + 1,
+            "content": match.group(1).strip(),
+            "raw": match.group(0),
+        })
+
+    # Display math with parenthesized numbers: $$ formula $$ (N)
+    display_math = re.compile(r"\$\$(.*?)\$\$\s*\((\d+)\)", re.DOTALL)
+    for match in display_math.finditer(text):
+        equations.append({
+            "number": int(match.group(2)),
+            "content": match.group(1).strip(),
+            "raw": match.group(0),
+        })
+
+    # Lines that look like equations with numbers at the end: formula (N)
+    numbered_line = re.compile(r"^(.+?)\s+\((\d+)\)\s*$", re.MULTILINE)
+    for match in numbered_line.finditer(text):
+        content = match.group(1).strip()
+        num = int(match.group(2))
+        # Only include if it looks like an equation (has math-like characters)
+        if any(c in content for c in "=+∑∏∫_^{}\\√∞"):
+            if not any(eq["number"] == num for eq in equations):
+                equations.append({
+                    "number": num,
+                    "content": content,
+                    "raw": match.group(0),
+                })
+
+    # Sort by equation number
+    equations.sort(key=lambda e: e["number"])
+    return equations
+
+
+def extract_tables(text: str) -> list[dict]:
+    """Extract tables from the paper text.
+
+    Looks for:
+      - Markdown tables (pipes)
+      - Table captions (Table N: ...)
+      - Structured text that looks like a table
+    """
+    tables = []
+
+    # Find table captions and associated content
+    table_caption = re.compile(
+        r"(Table\s+\d+[:\.]?\s*[^\n]*)\n(.*?)(?=Table\s+\d+[:\.]|^#{1,4}\s|\Z)",
+        re.DOTALL | re.MULTILINE,
+    )
+
+    for match in table_caption.finditer(text):
+        caption = match.group(1).strip()
+        body = match.group(2).strip()
+
+        # Check if the body contains table-like content (pipes, tabs, or aligned columns)
+        if "|" in body or "\t" in body or re.search(r"\s{3,}", body):
+            tables.append({
+                "caption": caption,
+                "content": body[:2000],  # limit size
+            })
+
+    # Also find markdown tables without explicit captions
+    md_table = re.compile(r"(\|[^\n]+\|\n\|[-:\s|]+\|\n(?:\|[^\n]+\|\n)*)", re.MULTILINE)
+    for match in md_table.finditer(text):
+        table_text = match.group(1).strip()
+        if not any(table_text in t["content"] for t in tables):
+            tables.append({
+                "caption": "Untitled table",
+                "content": table_text,
+            })
+
+    return tables
+
+
+def extract_footnotes(text: str) -> list[dict]:
+    """Extract footnotes from the paper."""
+    footnotes = []
+
+    # Pattern: footnote markers like ¹, ², ³ or [1], [2] at start of line
+    fn_pattern = re.compile(
+        r"(?:^|\n)[\s]*(?:[\u00b9\u00b2\u00b3\u2074-\u2079]|\[(\d+)\]|(\d+)\.)[\s]+(.+?)(?=\n[\s]*(?:[\u00b9\u00b2\u00b3\u2074-\u2079]|\[\d+\]|\d+\.)\s|\n\n|\Z)",
+        re.DOTALL,
+    )
+
+    for match in fn_pattern.finditer(text):
+        content = match.group(0).strip()
+        if len(content) > 10:  # skip very short matches that are likely false positives
+            footnotes.append(content)
+
+    # Also look for explicit footnote sections
+    fn_section = re.compile(
+        r"(?:footnote|note)s?\s*:?\s*\n(.*?)(?=\n#{1,4}\s|\Z)",
+        re.DOTALL | re.IGNORECASE,
+    )
+    for match in fn_section.finditer(text):
+        content = match.group(1).strip()
+        if content and content not in footnotes:
+            footnotes.append(content)
+
+    return [{"content": fn} for fn in footnotes]
+
+
+def save_list_to_dir(items: list[dict], output_dir: Path, name_key: str = "title"):
+    """Save a list of extracted items as individual files."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for i, item in enumerate(items):
+        # Create a clean filename
+        name = item.get(name_key, item.get("caption", f"item_{i+1}"))
+        name = str(name)
+        clean_name = re.sub(r"[^\w\s-]", "", name)
+        clean_name = re.sub(r"\s+", "_", clean_name).strip("_").lower()
+        if not clean_name:
+            clean_name = f"item_{i+1}"
+        clean_name = clean_name[:80]  # limit filename length
+
+        filepath = output_dir / f"{i+1:02d}_{clean_name}.md"
+
+        content = f"# {name}\n\n{item.get('content', item.get('raw', ''))}\n"
+        with open(filepath, "w", encoding="utf-8") as f:
+            f.write(content)
+
+
+def main():
+    if len(sys.argv) < 3:
+        print(f"Usage: {sys.argv[0]} <paper_text.md> <output_dir>", file=sys.stderr)
+        sys.exit(1)
+
+    paper_path = Path(sys.argv[1])
+    output_dir = Path(sys.argv[2])
+
+    if not paper_path.exists():
+        print(f"ERROR: {paper_path} does not exist", file=sys.stderr)
+        sys.exit(1)
+
+    print(f"Extracting structure from: {paper_path}")
+    text = paper_path.read_text(encoding="utf-8")
+    print(f"  Total characters: {len(text):,}")
+
+    # Extract sections
+    print("\n--- Extracting sections ---")
+    sections = identify_sections(text)
+    if sections:
+        save_list_to_dir(sections, output_dir / "sections")
+        print(f"  Found {len(sections)} sections:")
+        for s in sections:
+            print(f"    - {s['title']} ({len(s['content'])} chars)")
+    else:
+        print("  WARNING: No sections detected. The paper text may not have clear headings.")
+        # Save the entire text as a single section
+        (output_dir / "sections").mkdir(parents=True, exist_ok=True)
+        (output_dir / "sections" / "01_full_text.md").write_text(text, encoding="utf-8")
+
+    # Extract algorithms
+    print("\n--- Extracting algorithm boxes ---")
+    algorithms = extract_algorithms(text)
+    if algorithms:
+        save_list_to_dir(algorithms, output_dir / "algorithms")
+        print(f"  Found {len(algorithms)} algorithms:")
+        for a in algorithms:
+            print(f"    - {a['title']}")
+    else:
+        print("  No algorithm boxes found.")
+
+    # Extract equations
+    print("\n--- Extracting equations ---")
+    equations = extract_equations(text)
+    if equations:
+        save_list_to_dir(equations, output_dir / "equations", name_key="number")
+        print(f"  Found {len(equations)} numbered equations")
+    else:
+        print("  No numbered equations found (may be inline or in non-standard format).")
+
+    # Extract tables
+    print("\n--- Extracting tables ---")
+    tables = extract_tables(text)
+    if tables:
+        save_list_to_dir(tables, output_dir / "tables", name_key="caption")
+        print(f"  Found {len(tables)} tables:")
+        for t in tables:
+            print(f"    - {t['caption']}")
+    else:
+        print("  No tables found.")
+
+    # Extract footnotes
+    print("\n--- Extracting footnotes ---")
+    footnotes = extract_footnotes(text)
+    footnotes_path = output_dir / "footnotes.md"
+    if footnotes:
+        with open(footnotes_path, "w", encoding="utf-8") as f:
+            f.write("# Footnotes\n\n")
+            for i, fn in enumerate(footnotes):
+                f.write(f"## Footnote {i + 1}\n\n{fn['content']}\n\n---\n\n")
+        print(f"  Found {len(footnotes)} footnotes")
+    else:
+        with open(footnotes_path, "w", encoding="utf-8") as f:
+            f.write("# Footnotes\n\nNo footnotes extracted.\n")
+        print("  No footnotes found.")
+
+    # Summary
+    print(f"\n--- Extraction Summary ---")
+    print(f"  Sections:   {len(sections)}")
+    print(f"  Algorithms: {len(algorithms)}")
+    print(f"  Equations:  {len(equations)}")
+    print(f"  Tables:     {len(tables)}")
+    print(f"  Footnotes:  {len(footnotes)}")
+    print(f"  Output dir: {output_dir}")
+    print(f"\nDone.")
+
+
+if __name__ == "__main__":
+    main()

plugins/codex/plugins/paper2code/skills/paper2code/scripts/fetch_paper.py

@@ -0,0 +1,436 @@
+#!/usr/bin/env python3
+"""
+Fetch and parse an arxiv paper.
+
+Usage:
+    python fetch_paper.py <arxiv_id_or_url> <output_dir>
+
+Examples:
+    python fetch_paper.py 2106.09685 ./output/
+    python fetch_paper.py https://arxiv.org/abs/2106.09685 ./output/
+    python fetch_paper.py 2106.09685v2 ./output/
+
+Outputs:
+    {output_dir}/paper_text.md      — full paper text in markdown
+    {output_dir}/paper_metadata.json — title, authors, abstract, categories
+"""
+
+import json
+import re
+import sys
+from pathlib import Path
+
+import requests
+
+
+def normalize_arxiv_id(input_str: str) -> str:
+    """Extract arxiv ID from a URL or bare ID string.
+
+    Handles:
+        https://arxiv.org/abs/2106.09685
+        https://arxiv.org/pdf/2106.09685.pdf
+        http://arxiv.org/abs/2106.09685v2
+        2106.09685
+        2106.09685v2
+        cs/0601007  (old-style IDs)
+    """
+    input_str = input_str.strip().rstrip("/")
+
+    # Remove common URL prefixes
+    for prefix in [
+        "https://arxiv.org/abs/",
+        "http://arxiv.org/abs/",
+        "https://arxiv.org/pdf/",
+        "http://arxiv.org/pdf/",
+    ]:
+        if input_str.startswith(prefix):
+            input_str = input_str[len(prefix):]
+            break
+
+    # Remove .pdf suffix if present
+    if input_str.endswith(".pdf"):
+        input_str = input_str[:-4]
+
+    # Validate format: YYMM.NNNNN(vN) or archive/NNNNNNN
+    new_style = re.match(r"^\d{4}\.\d{4,5}(v\d+)?$", input_str)
+    old_style = re.match(r"^[a-z-]+/\d{7}(v\d+)?$", input_str)
+
+    if not new_style and not old_style:
+        print(f"WARNING: '{input_str}' may not be a valid arxiv ID.", file=sys.stderr)
+
+    return input_str
+
+
+def fetch_metadata(arxiv_id: str) -> dict:
+    """Fetch paper metadata from the arxiv API."""
+    # Strip version for API query
+    base_id = re.sub(r"v\d+$", "", arxiv_id)
+    api_url = f"http://export.arxiv.org/api/query?id_list={base_id}"
+
+    try:
+        resp = requests.get(api_url, timeout=30)
+        resp.raise_for_status()
+    except requests.RequestException as e:
+        print(f"WARNING: Could not fetch metadata from arxiv API: {e}", file=sys.stderr)
+        return {"arxiv_id": arxiv_id, "title": "Unknown", "authors": [], "abstract": "", "categories": []}
+
+    text = resp.text
+
+    # Simple XML parsing — avoid heavy dependencies
+    def extract_tag(tag: str, content: str) -> str:
+        pattern = rf"<{tag}[^>]*>(.*?)</{tag}>"
+        match = re.search(pattern, content, re.DOTALL)
+        return match.group(1).strip() if match else ""
+
+    def extract_all_tags(tag: str, content: str) -> list:
+        pattern = rf"<{tag}[^>]*>(.*?)</{tag}>"
+        return [m.strip() for m in re.findall(pattern, content, re.DOTALL)]
+
+    # Find the entry (skip the feed-level title)
+    entry_match = re.search(r"<entry>(.*?)</entry>", text, re.DOTALL)
+    if not entry_match:
+        print("WARNING: No entry found in arxiv API response.", file=sys.stderr)
+        return {"arxiv_id": arxiv_id, "title": "Unknown", "authors": [], "abstract": "", "categories": []}
+
+    entry = entry_match.group(1)
+
+    title = extract_tag("title", entry)
+    title = re.sub(r"\s+", " ", title)  # collapse whitespace
+
+    abstract = extract_tag("summary", entry)
+    abstract = re.sub(r"\s+", " ", abstract)
+
+    # Authors
+    author_names = []
+    for author_block in re.findall(r"<author>(.*?)</author>", entry, re.DOTALL):
+        name = extract_tag("name", author_block)
+        if name:
+            author_names.append(name)
+
+    # Categories
+    categories = re.findall(r'<category[^>]*term="([^"]+)"', entry)
+
+    return {
+        "arxiv_id": arxiv_id,
+        "title": title,
+        "authors": author_names,
+        "abstract": abstract,
+        "categories": categories,
+    }
+
+
+def download_pdf(arxiv_id: str, output_path: Path) -> bool:
+    """Download the PDF from arxiv."""
+    pdf_url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
+    print(f"Downloading PDF from {pdf_url}...")
+
+    try:
+        resp = requests.get(pdf_url, timeout=60, stream=True)
+        resp.raise_for_status()
+
+        with open(output_path, "wb") as f:
+            for chunk in resp.iter_content(chunk_size=8192):
+                f.write(chunk)
+
+        file_size = output_path.stat().st_size
+        print(f"  Downloaded: {file_size / 1024:.0f} KB")
+        return True
+
+    except requests.RequestException as e:
+        print(f"  FAILED: {e}", file=sys.stderr)
+        return False
+
+
+def extract_with_pymupdf4llm(pdf_path: Path) -> str | None:
+    """Extract text using pymupdf4llm (preserves math notation as LaTeX)."""
+    try:
+        import pymupdf4llm
+        print("Extracting with pymupdf4llm (math-preserving)...")
+        text = pymupdf4llm.to_markdown(str(pdf_path))
+        if text and len(text) > 500:
+            print(f"  Extracted: {len(text)} characters")
+            return text
+        print("  WARNING: pymupdf4llm produced insufficient text.", file=sys.stderr)
+        return None
+    except ImportError:
+        print("  pymupdf4llm not available.", file=sys.stderr)
+        return None
+    except Exception as e:
+        print(f"  pymupdf4llm failed: {e}", file=sys.stderr)
+        return None
+
+
+def extract_with_pdfplumber(pdf_path: Path) -> str | None:
+    """Extract text using pdfplumber (fallback)."""
+    try:
+        import pdfplumber
+        print("Extracting with pdfplumber (fallback)...")
+        pages = []
+        with pdfplumber.open(str(pdf_path)) as pdf:
+            for i, page in enumerate(pdf.pages):
+                text = page.extract_text()
+                if text:
+                    pages.append(f"<!-- Page {i + 1} -->\n{text}")
+        if pages:
+            full_text = "\n\n".join(pages)
+            print(f"  Extracted: {len(full_text)} characters from {len(pages)} pages")
+            return full_text
+        print("  WARNING: pdfplumber produced no text.", file=sys.stderr)
+        return None
+    except ImportError:
+        print("  pdfplumber not available.", file=sys.stderr)
+        return None
+    except Exception as e:
+        print(f"  pdfplumber failed: {e}", file=sys.stderr)
+        return None
+
+
+def fetch_ar5iv_html(arxiv_id: str) -> str | None:
+    """Fetch HTML version from ar5iv (renders math as readable text)."""
+    base_id = re.sub(r"v\d+$", "", arxiv_id)
+    html_url = f"https://ar5iv.labs.arxiv.org/html/{base_id}"
+    print(f"Fetching HTML from {html_url}...")
+
+    try:
+        resp = requests.get(html_url, timeout=60)
+        resp.raise_for_status()
+
+        # Basic HTML to text conversion — strip tags but keep structure
+        text = resp.text
+
+        # Remove script and style blocks
+        text = re.sub(r"<script[^>]*>.*?</script>", "", text, flags=re.DOTALL)
+        text = re.sub(r"<style[^>]*>.*?</style>", "", text, flags=re.DOTALL)
+
+        # Convert headers to markdown
+        for level in range(1, 7):
+            text = re.sub(
+                rf"<h{level}[^>]*>(.*?)</h{level}>",
+                lambda m, lv=level: f"\n{'#' * lv} {m.group(1).strip()}\n",
+                text,
+                flags=re.DOTALL,
+            )
+
+        # Convert paragraphs to double newlines
+        text = re.sub(r"<p[^>]*>", "\n\n", text)
+        text = re.sub(r"</p>", "", text)
+
+        # Convert list items
+        text = re.sub(r"<li[^>]*>", "\n- ", text)
+
+        # Preserve math elements (ar5iv uses MathML or LaTeX in alt text)
+        text = re.sub(r'<math[^>]*alttext="([^"]*)"[^>]*>.*?</math>', r"$\1$", text, flags=re.DOTALL)
+
+        # Strip remaining HTML tags
+        text = re.sub(r"<[^>]+>", "", text)
+
+        # Clean up whitespace
+        text = re.sub(r"\n{3,}", "\n\n", text)
+        text = text.strip()
+
+        if len(text) > 500:
+            print(f"  Extracted: {len(text)} characters from HTML")
+            return text
+
+        print("  WARNING: ar5iv HTML produced insufficient text.", file=sys.stderr)
+        return None
+
+    except requests.RequestException as e:
+        print(f"  ar5iv fetch failed: {e}", file=sys.stderr)
+        return None
+
+
+def check_text_quality(text: str) -> bool:
+    """Check if extracted text is reasonable quality (not garbled)."""
+    if not text or len(text) < 500:
+        return False
+
+    # Check first 1000 chars for readability
+    sample = text[:1000]
+
+    # Count non-ASCII, non-whitespace, non-LaTeX special chars
+    weird_chars = sum(
+        1 for c in sample
+        if ord(c) > 127 and c not in "αβγδεζηθικλμνξπρστυφχψωΓΔΘΛΞΠΣΦΨΩ∑∏∫∂∇√∞±≤≥≠≈∈∉⊂⊃∪∩"
+    )
+    weird_ratio = weird_chars / max(len(sample), 1)
+
+    if weird_ratio > 0.2:
+        print(f"  WARNING: Text quality check failed ({weird_ratio:.0%} non-standard characters)")
+        return False
+
+    # Check for recognizable English words
+    common_words = {"the", "and", "of", "in", "to", "we", "is", "for", "that", "with"}
+    words_lower = set(re.findall(r"\b[a-z]+\b", sample.lower()))
+    found_common = words_lower & common_words
+
+    if len(found_common) < 3:
+        print("  WARNING: Text quality check failed (few recognizable English words)")
+        return False
+
+    return True
+
+
+def find_official_code(arxiv_id: str, paper_text: str | None, metadata: dict) -> list[dict]:
+    """Search for official code repositories linked to this paper.
+
+    Checks two sources:
+    1. The paper text itself — GitHub/GitLab URLs, "code available at" phrases
+    2. The arxiv abstract page — authors sometimes add code links there
+
+    Returns a list of dicts with keys: url, source, context
+    """
+    found = []
+    seen_urls = set()
+
+    def add_link(url: str, source: str, context: str = "") -> None:
+        normalized = url.rstrip("/").lower()
+        if normalized not in seen_urls:
+            seen_urls.add(normalized)
+            found.append({"url": url.rstrip("/"), "source": source, "context": context.strip()})
+
+    # --- Source 1: Scan paper text for code URLs ---
+    if paper_text:
+        # Match GitHub/GitLab/Bitbucket repo URLs
+        repo_pattern = r"https?://(?:github\.com|gitlab\.com|bitbucket\.org)/[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+"
+        for match in re.finditer(repo_pattern, paper_text):
+            url = match.group(0)
+            # Grab surrounding context (up to 120 chars on each side)
+            start = max(0, match.start() - 120)
+            end = min(len(paper_text), match.end() + 120)
+            context = paper_text[start:end].replace("\n", " ")
+            add_link(url, "paper_text", context)
+
+        # Match common phrases that precede code URLs
+        code_phrases = [
+            r"code\s+(?:is\s+)?(?:available|released|open[\s-]?sourced)\s+at\s+(https?://\S+)",
+            r"(?:our|the)\s+code\s+(?:can be found|is hosted)\s+at\s+(https?://\S+)",
+            r"implementation\s+(?:is\s+)?(?:available|released)\s+at\s+(https?://\S+)",
+            r"source\s+code[:\s]+(https?://\S+)",
+        ]
+        for pattern in code_phrases:
+            for match in re.finditer(pattern, paper_text, re.IGNORECASE):
+                url = match.group(1).rstrip(".,;:)")
+                add_link(url, "paper_text", match.group(0))
+
+    # --- Source 2: Scan the arxiv abstract page ---
+    base_id = re.sub(r"v\d+$", "", arxiv_id)
+    abs_url = f"https://arxiv.org/abs/{base_id}"
+    try:
+        resp = requests.get(abs_url, timeout=30)
+        resp.raise_for_status()
+        html = resp.text
+
+        # arxiv shows official code links in the "Code" or "GitHub" badges / sidebar
+        # Look for GitHub links in the abstract page HTML
+        page_repo_matches = re.findall(
+            r'href="(https?://(?:github\.com|gitlab\.com|bitbucket\.org)/[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+)"',
+            html,
+        )
+        for url in page_repo_matches:
+            add_link(url, "arxiv_page", "Link found on arxiv abstract page")
+
+    except requests.RequestException as e:
+        print(f"  WARNING: Could not fetch arxiv abstract page for code links: {e}", file=sys.stderr)
+
+    return found
+
+
+def main():
+    if len(sys.argv) < 3:
+        print(f"Usage: {sys.argv[0]} <arxiv_id_or_url> <output_dir>", file=sys.stderr)
+        sys.exit(1)
+
+    raw_input = sys.argv[1]
+    output_dir = Path(sys.argv[2])
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Step 1: Normalize ID
+    arxiv_id = normalize_arxiv_id(raw_input)
+    print(f"Arxiv ID: {arxiv_id}")
+
+    # Step 2: Fetch metadata
+    print("\n--- Fetching metadata ---")
+    metadata = fetch_metadata(arxiv_id)
+    metadata_path = output_dir / "paper_metadata.json"
+    with open(metadata_path, "w", encoding="utf-8") as f:
+        json.dump(metadata, f, indent=2, ensure_ascii=False)
+    print(f"  Title: {metadata['title']}")
+    print(f"  Authors: {', '.join(metadata['authors'][:5])}{'...' if len(metadata['authors']) > 5 else ''}")
+    print(f"  Categories: {', '.join(metadata['categories'])}")
+
+    # Step 3: Download and extract PDF
+    paper_text = None
+    pdf_path = output_dir / "paper.pdf"
+
+    print("\n--- Downloading PDF ---")
+    if download_pdf(arxiv_id, pdf_path):
+        # Try pymupdf4llm first
+        print("\n--- Extracting text ---")
+        paper_text = extract_with_pymupdf4llm(pdf_path)
+
+        # Check quality
+        if paper_text and not check_text_quality(paper_text):
+            print("  pymupdf4llm text quality poor, trying pdfplumber...")
+            paper_text = None
+
+        # Fallback to pdfplumber
+        if paper_text is None:
+            paper_text = extract_with_pdfplumber(pdf_path)
+
+        if paper_text and not check_text_quality(paper_text):
+            print("  pdfplumber text quality poor, trying ar5iv HTML...")
+            paper_text = None
+
+    # Step 4: Fallback to ar5iv HTML
+    if paper_text is None:
+        print("\n--- Trying ar5iv HTML fallback ---")
+        paper_text = fetch_ar5iv_html(arxiv_id)
+
+    # Step 5: Save results
+    if paper_text is None:
+        print("\nERROR: All extraction methods failed.", file=sys.stderr)
+        print("Please download the paper manually and provide the text.", file=sys.stderr)
+        sys.exit(1)
+
+    text_path = output_dir / "paper_text.md"
+    with open(text_path, "w", encoding="utf-8") as f:
+        f.write(f"# {metadata['title']}\n\n")
+        f.write(f"**Authors:** {', '.join(metadata['authors'])}\n\n")
+        f.write(f"**ArXiv:** https://arxiv.org/abs/{arxiv_id}\n\n")
+        f.write("---\n\n")
+        f.write(paper_text)
+
+    # Step 6: Search for official code repositories
+    code_links = find_official_code(arxiv_id, paper_text, metadata)
+    if code_links:
+        metadata["official_code"] = code_links
+        # Re-save metadata with code links
+        with open(metadata_path, "w", encoding="utf-8") as f:
+            json.dump(metadata, f, indent=2, ensure_ascii=False)
+        for link in code_links:
+            print(f"  Found: {link['url']} (source: {link['source']})")
+    else:
+        metadata["official_code"] = []
+        with open(metadata_path, "w", encoding="utf-8") as f:
+            json.dump(metadata, f, indent=2, ensure_ascii=False)
+        print("  No official code repositories found.")
+
+    # Summary
+    page_count = paper_text.count("<!-- Page")
+    has_math = bool(re.search(r"[\$\\]|\\frac|\\sum|\\int|\\mathbb", paper_text))
+    has_figures = bool(re.search(r"[Ff]igure\s+\d", paper_text))
+
+    print(f"\n--- Extraction Summary ---")
+    print(f"  Output: {text_path}")
+    print(f"  Characters: {len(paper_text):,}")
+    print(f"  Pages detected: {page_count if page_count > 0 else 'N/A (HTML extraction)'}")
+    print(f"  Math preserved: {'Yes' if has_math else 'No'}")
+    print(f"  Figure references found: {'Yes' if has_figures else 'No'}")
+    print(f"  Metadata saved: {metadata_path}")
+    print(f"  Official code links: {len(code_links)} found")
+    print(f"\nDone.")
+
+
+if __name__ == "__main__":
+    main()

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/README.md

@@ -0,0 +1,69 @@
+# Attention Is All You Need
+
+Implementation of [Attention Is All You Need](https://arxiv.org/abs/1706.03762) (Vaswani et al., 2017).
+
+## What this implements
+
+The Transformer architecture — an encoder-decoder model based entirely on multi-head self-attention mechanisms, eliminating recurrence and convolutions. This implementation covers the full architecture as described in §3 of the paper: scaled dot-product attention, multi-head attention, positional encoding, the encoder and decoder stacks, and the output projection with weight tying.
+
+## Quick start
+
+```bash
+pip install -r requirements.txt
+```
+
+```python
+from src.model import Transformer, TransformerConfig
+
+config = TransformerConfig()
+model = Transformer(config)
+
+# Example: encoder-decoder forward pass
+import torch
+src = torch.randint(0, config.vocab_size, (2, 20))   # (batch, src_len)
+tgt = torch.randint(0, config.vocab_size, (2, 15))   # (batch, tgt_len)
+output = model(src, tgt)
+print(output.shape)  # (2, 15, vocab_size)
+```
+
+## File structure
+
+```
+attention_is_all_you_need/
+├── README.md                 # This file
+├── REPRODUCTION_NOTES.md     # Ambiguity audit — what's specified vs. assumed
+├── requirements.txt          # Dependencies
+├── src/
+│   ├── model.py              # Full Transformer architecture (§3)
+│   ├── loss.py               # Label-smoothed cross-entropy (§5.4)
+│   ├── data.py               # Dataset skeleton for WMT translation
+│   ├── train.py              # Minimal training example
+│   ├── evaluate.py           # BLEU score computation
+│   └── utils.py              # Masking utilities, LR schedule
+├── configs/
+│   └── base.yaml             # All hyperparameters from §3 and Table 3
+└── notebooks/
+    └── walkthrough.ipynb     # Paper sections → code → sanity checks
+```
+
+## Important: Read REPRODUCTION_NOTES.md
+
+This implementation flags every choice that the paper does not specify.
+Before using this code for research, read [REPRODUCTION_NOTES.md](REPRODUCTION_NOTES.md)
+to understand which implementation details are from the paper and which are our choices.
+
+## Citation
+
+```bibtex
+@article{vaswani2017attention,
+  title={Attention is all you need},
+  author={Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and Kaiser, {\L}ukasz and Polosukhin, Illia},
+  journal={Advances in neural information processing systems},
+  volume={30},
+  year={2017}
+}
+```
+
+---
+
+*Generated by [paper2code](https://github.com/PrathamLearnsToCode/paper2code) — citation-anchored paper implementation.*

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/REPRODUCTION_NOTES.md

@@ -0,0 +1,118 @@
+# Reproduction Notes: Attention Is All You Need
+
+> This document records every implementation choice, whether it was specified by the paper,
+> and what alternatives exist. If you're reproducing this paper, **read this first.**
+
+---
+
+## Paper
+
+- **Title:** Attention Is All You Need
+- **Authors:** Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin
+- **Year:** 2017
+- **ArXiv:** https://arxiv.org/abs/1706.03762
+- **Official code:** https://github.com/tensorflow/tensor2tensor
+
+---
+
+## What this implements
+
+The Transformer model architecture — an encoder-decoder sequence-to-sequence model that replaces recurrence entirely with multi-head self-attention. The encoder maps an input sequence of symbol representations to a sequence of continuous representations, and the decoder generates an output sequence one element at a time using attention over the encoder output and previously generated symbols. This is the "base" model configuration (§3, Table 3).
+
+---
+
+## Verified against
+
+- [x] Paper equations (§3.2.1 Eq. 1 — scaled dot-product attention, §3.2.2 — multi-head attention, §3.5 — positional encoding)
+- [x] Paper Table 3 — architecture hyperparameters
+- [x] Paper §5.3 — optimizer configuration
+- [ ] Official code (tensor2tensor) — referenced but not line-by-line verified
+- [x] Well-known reimplementations: Harvard NLP "The Annotated Transformer"
+
+---
+
+## Unspecified choices
+
+| Component | Our Choice | Alternatives | Paper Quote (if partial) | Section |
+|-----------|-----------|--------------|--------------------------|---------|
+| LayerNorm epsilon | 1e-6 | 1e-5 (PyTorch default), 1e-8 | — | — |
+| LayerNorm placement | Post-norm | Pre-norm (more stable at scale) | "add & norm" in Figure 1 | §3.1, Figure 1 |
+| Activation in FFN | ReLU | GELU (modern default) | "ReLU" in Eq. 2 | §3.3 |
+| Bias in linear projections | True | False (used in some reimplementations) | — | — |
+| Bias in output projection | True | False | — | — |
+| Weight initialization | Xavier uniform | PyTorch defaults, normal(0, 0.02) | — | — |
+| Embedding scale factor | √d_model | 1.0 (no scaling) | "multiply those weights by √d_model" | §3.4 |
+| Weight tying | Embedding + output projection shared | Separate weights | "share the same weight matrix" | §3.4 |
+| Dropout on attention weights | Yes (rate 0.1) | No dropout on weights | — | §5.4 states P_drop=0.1 but placement unspecified |
+| Dropout placement | After attention + after FFN + on embeddings | Various | "Residual Dropout... applied to output of each sub-layer" | §5.4 |
+| Max sequence length for PE | 5000 | 512, 10000 | — | — |
+
+---
+
+## Known deviations
+
+| Deviation | Paper says | We do | Reason |
+|-----------|-----------|-------|--------|
+| None intentional | — | — | Implementation follows paper specification closely |
+
+---
+
+## Expected results
+
+| Metric | Paper's number | Dataset | Conditions |
+|--------|---------------|---------|------------|
+| BLEU | 27.3 | WMT 2014 EN-DE | Table 2, base model, beam search with beam size 4 |
+| BLEU | 38.1 | WMT 2014 EN-FR | Table 2, big model |
+| Training cost | 12 hours on 8 P100 GPUs | — | §5.2, base model, 100k steps |
+
+**Note:** Exact reproduction requires the WMT 2014 dataset with the same preprocessing
+(BPE tokenization with ~37k tokens for EN-DE), 8 P100 GPUs, and the exact optimizer
+schedule. Small deviations in BLEU (0.1-0.5) are normal.
+
+---
+
+## Debugging tips
+
+1. **Loss not decreasing**: Check that the learning rate schedule follows Eq. 3 exactly (`lr = d_model^(-0.5) * min(step^(-0.5), step * warmup^(-1.5))`). Using a constant LR or a generic warmup schedule will give different results.
+
+2. **NaN in attention**: Check attention masking. The mask should add `-inf` (or a very large negative number) BEFORE softmax, not multiply AFTER. All-zero rows in the mask (fully masked positions) can cause NaN after softmax — ensure at least one position is always attended to.
+
+3. **Poor translation quality**: Verify (a) label smoothing is ε=0.1 with KL divergence, (b) the embedding scaling factor √d_model is applied, (c) weight tying between embedding and output projection is enabled.
+
+4. **Slow convergence**: The warmup schedule is critical. With 4000 warmup steps, the LR peaks around step 4000 then decays. Starting with a high LR (no warmup) often diverges.
+
+5. **Parameter count mismatch**: The base model should have approximately 65M parameters. If significantly different, check weight tying and whether the vocabulary size matches.
+
+---
+
+## Scope decisions
+
+### Implemented
+- Scaled dot-product attention (§3.2.1) — core contribution
+- Multi-head attention (§3.2.2) — core contribution
+- Positional encoding, sinusoidal (§3.5) — core contribution
+- Encoder stack with N=6 layers (§3.1) — core contribution
+- Decoder stack with N=6 layers (§3.1) — core contribution
+- Embedding with √d_model scaling and weight tying (§3.4) — specified in paper
+- Label-smoothed cross-entropy loss (§5.4) — needed for reproduction
+
+### Intentionally excluded
+- BPE tokenization — standard preprocessing, use sentencepiece or HuggingFace tokenizers
+- Beam search decoding — inference technique, not the architectural contribution
+- Multi-GPU training — engineering concern, not the paper's contribution
+- WMT dataset downloading — data acquisition, provide instructions only
+
+### Needed for full reproduction (not included)
+- WMT 2014 EN-DE dataset with BPE tokenization (~37k merge operations)
+- 8 P100 GPUs for the reported training speed
+- Beam search with beam size 4 and length penalty α=0.6 for Table 2 results
+
+---
+
+## References
+
+- Vaswani et al., 2017 — primary paper
+- Ba et al., 2016 (Layer Normalization) — used in each sub-layer
+- Srivastava et al., 2014 (Dropout) — applied at P_drop=0.1
+- Szegedy et al., 2016 (Label smoothing) — ε_ls=0.1
+- The Annotated Transformer (Harvard NLP) — reference reimplementation

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/configs/base.yaml

@@ -0,0 +1,59 @@
+# Configuration for Attention Is All You Need (Vaswani et al., 2017)
+# Paper: https://arxiv.org/abs/1706.03762
+#
+# CONVENTION:
+#   - Lines with §X.Y cite the paper section specifying this value
+#   - Lines with [UNSPECIFIED] mark choices not stated in the paper
+#   - Lines with [FROM_TABLE] reference Table 3 (Architecture hyperparameters)
+#
+# All hyperparameters in one place. No magic numbers in code.
+
+# ---------------------------------------------------------------------------
+# Model architecture — §3, Table 3 ("base model")
+# ---------------------------------------------------------------------------
+model:
+  d_model: 512             # §3.1, Table 3 — "d_model = 512"
+  n_heads: 8               # §3.2.2, Table 3 — "h = 8"
+  d_ff: 2048               # §3.3, Table 3 — "d_ff = 2048"
+  n_encoder_layers: 6      # §3.1, Table 3 — "N = 6" (encoder)
+  n_decoder_layers: 6      # §3.1, Table 3 — "N = 6" (decoder)
+  dropout: 0.1             # §5.4, Table 3 — "P_drop = 0.1"
+  max_seq_len: 5000        # [UNSPECIFIED] — max length for positional encoding
+  vocab_size: 37000        # §5.1 — "byte-pair encoding... shared source-target vocabulary of about 37000 tokens" (EN-DE)
+  activation: "relu"       # §3.3, Eq. 2 — "max(0, xW_1 + b_1)W_2 + b_2" (ReLU)
+  norm_eps: 1.0e-6         # [UNSPECIFIED] — LayerNorm epsilon, paper does not state
+  tie_weights: true        # §3.4 — "we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation"
+
+# ---------------------------------------------------------------------------
+# Training — §5.3, §5.4
+# ---------------------------------------------------------------------------
+training:
+  optimizer: "adam"         # §5.3 — "We used the Adam optimizer"
+  betas: [0.9, 0.98]       # §5.3 — "β1 = 0.9, β2 = 0.98"
+  eps: 1.0e-9              # §5.3 — "ε = 10^−9"
+  label_smoothing: 0.1     # §5.4 — "label smoothing of value ε_ls = 0.1"
+  warmup_steps: 4000       # §5.3 — "warmup_steps = 4000"
+  total_steps: 100000      # §5.2 — "base models were trained for a total of 100,000 steps"
+  batch_tokens: 25000      # §5.1 — "Each training batch contained a set of sentence pairs containing approximately 25000 source tokens and 25000 target tokens"
+  weight_decay: 0.0        # [UNSPECIFIED] — no weight decay mentioned
+  gradient_clip: null      # [UNSPECIFIED] — no gradient clipping mentioned
+
+  # LR schedule — §5.3, Eq. 3
+  # lr = d_model^(-0.5) * min(step_num^(-0.5), step_num * warmup_steps^(-1.5))
+  schedule: "transformer"  # §5.3, Eq. 3 — custom schedule, not cosine or linear
+
+# ---------------------------------------------------------------------------
+# Data — §5.1
+# ---------------------------------------------------------------------------
+data:
+  dataset: "WMT 2014 EN-DE"       # §5.1 — "WMT 2014 English-German dataset"
+  bpe_tokens: 37000                # §5.1 — "byte-pair encoding... about 37000 tokens"
+  max_seq_len: 512                 # [UNSPECIFIED] — max sequence length for training
+
+# ---------------------------------------------------------------------------
+# Evaluation — §6
+# ---------------------------------------------------------------------------
+eval:
+  metric: "BLEU"                   # §6 — "BLEU score"
+  beam_size: 4                     # §6.1 — "beam search with a beam size of 4"
+  length_penalty: 0.6              # §6.1 — "length penalty α = 0.6"

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/notebooks/walkthrough.ipynb

@@ -0,0 +1,658 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e5efb596",
+   "metadata": {},
+   "source": [
+    "# Attention Is All You Need — Walkthrough\n",
+    "\n",
+    "**Paper:** [Attention Is All You Need](https://arxiv.org/abs/1706.03762) (Vaswani et al., 2017)\n",
+    "\n",
+    "This notebook walks through the Transformer architecture, connecting each paper section to code.\n",
+    "All examples use **tiny dimensions** so everything runs instantly on CPU.\n",
+    "The architecture is identical to the paper — only the numbers change.\n",
+    "\n",
+    "For each component:\n",
+    "1. Paper excerpt (the relevant passage)\n",
+    "2. Code (the implementation)\n",
+    "3. Sanity check (verify shapes and behavior)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9c8a4651",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "from dataclasses import dataclass\n",
+    "from typing import Optional\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "# Toy dimensions — paper uses d_model=512, we use 64 for CPU walkthrough\n",
+    "BATCH_SIZE = 2\n",
+    "SRC_LEN = 10\n",
+    "TGT_LEN = 8\n",
+    "D_MODEL = 64      # Paper: 512\n",
+    "N_HEADS = 4        # Paper: 8\n",
+    "D_FF = 128         # Paper: 2048\n",
+    "N_LAYERS = 2       # Paper: 6\n",
+    "VOCAB_SIZE = 100   # Paper: 37000\n",
+    "DROPOUT = 0.0      # Set to 0 for deterministic walkthrough\n",
+    "\n",
+    "print(f\"Using toy dimensions: d_model={D_MODEL}, n_heads={N_HEADS}, d_ff={D_FF}\")\n",
+    "print(f\"d_k = d_model / n_heads = {D_MODEL // N_HEADS}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c8346be9",
+   "metadata": {},
+   "source": [
+    "## 1. Scaled Dot-Product Attention (§3.2.1)\n",
+    "\n",
+    "> \"An attention function can be described as mapping a query and a set of key-value\n",
+    "> pairs to an output, where the query, keys, values, and output are all vectors.\n",
+    "> The output is computed as a weighted sum of the values, where the weight assigned\n",
+    "> to each value is computed by a compatibility function of the query with the\n",
+    "> corresponding key.\"\n",
+    ">\n",
+    "> — §3.2, Attention Is All You Need\n",
+    "\n",
+    "**Equation 1:**\n",
+    "$$\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V$$"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a40f5933",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def scaled_dot_product_attention(\n",
+    "    query: torch.Tensor,\n",
+    "    key: torch.Tensor,\n",
+    "    value: torch.Tensor,\n",
+    "    mask: Optional[torch.Tensor] = None,\n",
+    ") -> torch.Tensor:\n",
+    "    \"\"\"§3.2.1, Eq. 1 — Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V\"\"\"\n",
+    "    d_k = query.size(-1)\n",
+    "    \n",
+    "    # (batch, n_heads, seq_q, d_k) @ (batch, n_heads, d_k, seq_k) -> (batch, n_heads, seq_q, seq_k)\n",
+    "    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)\n",
+    "    \n",
+    "    if mask is not None:\n",
+    "        scores = scores.masked_fill(mask == 0, float('-inf'))\n",
+    "    \n",
+    "    attn_weights = F.softmax(scores, dim=-1)  # (batch, n_heads, seq_q, seq_k)\n",
+    "    return torch.matmul(attn_weights, value)   # (batch, n_heads, seq_q, d_v)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ee66cd76",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Sanity check: scaled dot-product attention\n",
+    "d_k = D_MODEL // N_HEADS\n",
+    "q = torch.randn(BATCH_SIZE, N_HEADS, SRC_LEN, d_k)\n",
+    "k = torch.randn(BATCH_SIZE, N_HEADS, SRC_LEN, d_k)\n",
+    "v = torch.randn(BATCH_SIZE, N_HEADS, SRC_LEN, d_k)\n",
+    "\n",
+    "attn_out = scaled_dot_product_attention(q, k, v)\n",
+    "assert attn_out.shape == (BATCH_SIZE, N_HEADS, SRC_LEN, d_k), \\\n",
+    "    f\"Expected {(BATCH_SIZE, N_HEADS, SRC_LEN, d_k)}, got {attn_out.shape}\"\n",
+    "print(f\"✓ Scaled dot-product attention: ({BATCH_SIZE}, {N_HEADS}, {SRC_LEN}, {d_k}) -> {attn_out.shape}\")\n",
+    "\n",
+    "# Verify attention weights sum to 1\n",
+    "scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)\n",
+    "weights = F.softmax(scores, dim=-1)\n",
+    "weight_sums = weights.sum(dim=-1)\n",
+    "assert torch.allclose(weight_sums, torch.ones_like(weight_sums), atol=1e-6), \\\n",
+    "    \"Attention weights should sum to 1\"\n",
+    "print(f\"✓ Attention weights sum to 1 (verified)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9585614f",
+   "metadata": {},
+   "source": [
+    "## 2. Multi-Head Attention (§3.2.2)\n",
+    "\n",
+    "> \"Multi-head attention allows the model to jointly attend to information from\n",
+    "> different representation subspaces at different positions. With a single attention\n",
+    "> head, averaging inhibits this.\"\n",
+    ">\n",
+    "> \"MultiHead(Q, K, V) = Concat(head₁, ..., headₕ)W^O\n",
+    ">  where headᵢ = Attention(QWᵢ^Q, KWᵢ^K, VWᵢ^V)\"\n",
+    ">\n",
+    "> \"We employ h = 8 parallel attention layers, or heads. For each of these we use\n",
+    "> d_k = d_v = d_model/h = 64.\"\n",
+    ">\n",
+    "> — §3.2.2, Attention Is All You Need"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9ab01cf5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MultiHeadAttention(nn.Module):\n",
+    "    \"\"\"§3.2.2 — Multi-Head Attention.\"\"\"\n",
+    "    \n",
+    "    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0):\n",
+    "        super().__init__()\n",
+    "        assert d_model % n_heads == 0\n",
+    "        self.d_model = d_model\n",
+    "        self.n_heads = n_heads\n",
+    "        self.d_k = d_model // n_heads  # §3.2.2 — d_k = d_model/h\n",
+    "        \n",
+    "        # §3.2.2 — projection matrices W_Q, W_K, W_V, W_O\n",
+    "        self.W_q = nn.Linear(d_model, d_model)\n",
+    "        self.W_k = nn.Linear(d_model, d_model)\n",
+    "        self.W_v = nn.Linear(d_model, d_model)\n",
+    "        self.W_o = nn.Linear(d_model, d_model)\n",
+    "        self.dropout = nn.Dropout(dropout)\n",
+    "    \n",
+    "    def forward(self, query, key, value, mask=None):\n",
+    "        batch_size = query.size(0)\n",
+    "        \n",
+    "        # Project and reshape: (batch, seq, d_model) -> (batch, n_heads, seq, d_k)\n",
+    "        q = self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)\n",
+    "        k = self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)\n",
+    "        v = self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)\n",
+    "        \n",
+    "        # Scaled dot-product attention\n",
+    "        attn_out = scaled_dot_product_attention(q, k, v, mask=mask)\n",
+    "        \n",
+    "        # Concatenate and project: (batch, n_heads, seq, d_k) -> (batch, seq, d_model)\n",
+    "        attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)\n",
+    "        return self.W_o(attn_out)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "61e16c15",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Sanity check: multi-head attention\n",
+    "mha = MultiHeadAttention(D_MODEL, N_HEADS, dropout=0.0)\n",
+    "x = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
+    "out = mha(x, x, x)  # self-attention\n",
+    "assert out.shape == (BATCH_SIZE, SRC_LEN, D_MODEL), \\\n",
+    "    f\"Expected {(BATCH_SIZE, SRC_LEN, D_MODEL)}, got {out.shape}\"\n",
+    "print(f\"✓ MultiHeadAttention: (batch={BATCH_SIZE}, seq={SRC_LEN}, d_model={D_MODEL}) -> {out.shape}\")\n",
+    "print(f\"  Parameters: {sum(p.numel() for p in mha.parameters()):,}\")\n",
+    "print(f\"  Expected: 4 * d_model^2 + 4 * d_model (biases) = {4 * D_MODEL**2 + 4 * D_MODEL:,}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f6191f70",
+   "metadata": {},
+   "source": [
+    "## 3. Position-wise Feed-Forward Network (§3.3)\n",
+    "\n",
+    "> \"In addition to attention sub-layers, each of the layers in our encoder and decoder\n",
+    "> contains a fully connected feed-forward network, which is applied to each position\n",
+    "> separately and identically.\"\n",
+    ">\n",
+    "> \"FFN(x) = max(0, xW₁ + b₁)W₂ + b₂\" (Eq. 2)\n",
+    ">\n",
+    "> \"The dimensionality of input and output is d_model = 512, and the inner-layer has\n",
+    "> dimensionality d_ff = 2048.\"\n",
+    ">\n",
+    "> — §3.3, Attention Is All You Need"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a52d14af",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FeedForward(nn.Module):\n",
+    "    \"\"\"§3.3, Eq. 2 — FFN(x) = max(0, xW_1 + b_1)W_2 + b_2\"\"\"\n",
+    "    \n",
+    "    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0):\n",
+    "        super().__init__()\n",
+    "        self.linear1 = nn.Linear(d_model, d_ff)\n",
+    "        self.linear2 = nn.Linear(d_ff, d_model)\n",
+    "        self.dropout = nn.Dropout(dropout)\n",
+    "    \n",
+    "    def forward(self, x):\n",
+    "        # §3.3, Eq. 2 — ReLU activation\n",
+    "        x = self.linear1(x)     # (batch, seq, d_model) -> (batch, seq, d_ff)\n",
+    "        x = F.relu(x)           # max(0, ...)\n",
+    "        x = self.dropout(x)\n",
+    "        x = self.linear2(x)     # (batch, seq, d_ff) -> (batch, seq, d_model)\n",
+    "        return x\n",
+    "\n",
+    "# Sanity check\n",
+    "ff = FeedForward(D_MODEL, D_FF)\n",
+    "x = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
+    "out = ff(x)\n",
+    "assert out.shape == (BATCH_SIZE, SRC_LEN, D_MODEL)\n",
+    "print(f\"✓ FeedForward: {x.shape} -> {out.shape}\")\n",
+    "print(f\"  Inner dimension: {D_FF} (expansion factor: {D_FF/D_MODEL}x)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "92ca2960",
+   "metadata": {},
+   "source": [
+    "## 4. Positional Encoding (§3.5)\n",
+    "\n",
+    "> \"Since our model contains no recurrence and no convolution, in order for the model\n",
+    "> to make use of the order of the sequence, we must inject some information about the\n",
+    "> relative or absolute position of the tokens in the sequence.\"\n",
+    ">\n",
+    "> \"We use sine and cosine functions of different frequencies:\n",
+    ">   PE(pos, 2i) = sin(pos / 10000^(2i/d_model))\n",
+    ">   PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))\"\n",
+    ">\n",
+    "> — §3.5, Attention Is All You Need"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "63a07ce6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class PositionalEncoding(nn.Module):\n",
+    "    \"\"\"§3.5 — Sinusoidal positional encoding.\"\"\"\n",
+    "    \n",
+    "    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.0):\n",
+    "        super().__init__()\n",
+    "        self.dropout = nn.Dropout(dropout)\n",
+    "        \n",
+    "        pe = torch.zeros(max_len, d_model)\n",
+    "        position = torch.arange(0, max_len).unsqueeze(1).float()\n",
+    "        div_term = torch.exp(\n",
+    "            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)\n",
+    "        )\n",
+    "        pe[:, 0::2] = torch.sin(position * div_term)  # even: sin\n",
+    "        pe[:, 1::2] = torch.cos(position * div_term)  # odd: cos\n",
+    "        pe = pe.unsqueeze(0)  # (1, max_len, d_model)\n",
+    "        self.register_buffer('pe', pe)\n",
+    "    \n",
+    "    def forward(self, x):\n",
+    "        x = x + self.pe[:, :x.size(1)]\n",
+    "        return self.dropout(x)\n",
+    "\n",
+    "# Sanity check\n",
+    "pos_enc = PositionalEncoding(D_MODEL, max_len=100)\n",
+    "x = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
+    "out = pos_enc(x)\n",
+    "assert out.shape == x.shape\n",
+    "print(f\"✓ PositionalEncoding: {x.shape} -> {out.shape} (added, not concatenated)\")\n",
+    "\n",
+    "# Verify PE values are bounded (sin/cos are in [-1, 1])\n",
+    "pe_values = pos_enc.pe[0, :SRC_LEN, :]\n",
+    "assert pe_values.min() >= -1.0 and pe_values.max() <= 1.0\n",
+    "print(f\"✓ PE values bounded in [-1, 1]: min={pe_values.min():.3f}, max={pe_values.max():.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "067a45be",
+   "metadata": {},
+   "source": [
+    "## 5. Encoder Layer (§3.1)\n",
+    "\n",
+    "> \"The encoder is composed of a stack of N = 6 identical layers. Each layer has two\n",
+    "> sub-layers. The first is a multi-head self-attention mechanism, and the second is a\n",
+    "> simple, position-wise fully connected feed-forward network.\"\n",
+    ">\n",
+    "> \"We employ a residual connection around each of the two sub-layers, followed by\n",
+    "> layer normalization. That is, the output of each sub-layer is LayerNorm(x + Sublayer(x)).\"\n",
+    ">\n",
+    "> — §3.1, Attention Is All You Need"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6a7aa6d3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class EncoderLayer(nn.Module):\n",
+    "    \"\"\"§3.1 — Single encoder layer: self-attention + feed-forward, with residual + norm.\"\"\"\n",
+    "    \n",
+    "    def __init__(self, d_model, n_heads, d_ff, dropout=0.0, norm_eps=1e-6):\n",
+    "        super().__init__()\n",
+    "        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)\n",
+    "        self.feed_forward = FeedForward(d_model, d_ff, dropout)\n",
+    "        # [UNSPECIFIED] LayerNorm epsilon not stated in paper — using 1e-6\n",
+    "        self.norm1 = nn.LayerNorm(d_model, eps=norm_eps)\n",
+    "        self.norm2 = nn.LayerNorm(d_model, eps=norm_eps)\n",
+    "        self.dropout1 = nn.Dropout(dropout)\n",
+    "        self.dropout2 = nn.Dropout(dropout)\n",
+    "    \n",
+    "    def forward(self, x, mask=None):\n",
+    "        # §3.1 — Post-norm: LayerNorm(x + Sublayer(x))\n",
+    "        attn_out = self.self_attn(x, x, x, mask)  # self-attention\n",
+    "        x = self.norm1(x + self.dropout1(attn_out))\n",
+    "        ff_out = self.feed_forward(x)\n",
+    "        x = self.norm2(x + self.dropout2(ff_out))\n",
+    "        return x\n",
+    "\n",
+    "# Sanity check\n",
+    "enc_layer = EncoderLayer(D_MODEL, N_HEADS, D_FF)\n",
+    "x = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
+    "out = enc_layer(x)\n",
+    "assert out.shape == (BATCH_SIZE, SRC_LEN, D_MODEL)\n",
+    "print(f\"✓ EncoderLayer: {x.shape} -> {out.shape}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0b658381",
+   "metadata": {},
+   "source": [
+    "## 6. Decoder Layer (§3.1)\n",
+    "\n",
+    "> \"In addition to the two sub-layers in each encoder layer, the decoder inserts a\n",
+    "> third sub-layer, which performs multi-head attention over the output of the encoder\n",
+    "> stack.\"\n",
+    ">\n",
+    "> \"We also modify the self-attention sub-layer in the decoder stack to prevent\n",
+    "> positions from attending to subsequent positions. This masking, combined with the\n",
+    "> fact that the output embeddings are offset by one position, ensures that the\n",
+    "> predictions for position i can depend only on the known outputs at positions less\n",
+    "> than i.\"\n",
+    ">\n",
+    "> — §3.1, §3.2.3, Attention Is All You Need"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ffdaf93e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class DecoderLayer(nn.Module):\n",
+    "    \"\"\"§3.1 — Single decoder layer: masked self-attn + cross-attn + FFN.\"\"\"\n",
+    "    \n",
+    "    def __init__(self, d_model, n_heads, d_ff, dropout=0.0, norm_eps=1e-6):\n",
+    "        super().__init__()\n",
+    "        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)\n",
+    "        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)\n",
+    "        self.feed_forward = FeedForward(d_model, d_ff, dropout)\n",
+    "        self.norm1 = nn.LayerNorm(d_model, eps=norm_eps)\n",
+    "        self.norm2 = nn.LayerNorm(d_model, eps=norm_eps)\n",
+    "        self.norm3 = nn.LayerNorm(d_model, eps=norm_eps)\n",
+    "        self.dropout1 = nn.Dropout(dropout)\n",
+    "        self.dropout2 = nn.Dropout(dropout)\n",
+    "        self.dropout3 = nn.Dropout(dropout)\n",
+    "    \n",
+    "    def forward(self, x, memory, tgt_mask=None, memory_mask=None):\n",
+    "        # Sub-layer 1: Masked self-attention\n",
+    "        self_attn_out = self.self_attn(x, x, x, tgt_mask)\n",
+    "        x = self.norm1(x + self.dropout1(self_attn_out))\n",
+    "        # Sub-layer 2: Cross-attention over encoder output\n",
+    "        cross_attn_out = self.cross_attn(x, memory, memory, memory_mask)\n",
+    "        x = self.norm2(x + self.dropout2(cross_attn_out))\n",
+    "        # Sub-layer 3: Feed-forward\n",
+    "        ff_out = self.feed_forward(x)\n",
+    "        x = self.norm3(x + self.dropout3(ff_out))\n",
+    "        return x\n",
+    "\n",
+    "# Sanity check\n",
+    "dec_layer = DecoderLayer(D_MODEL, N_HEADS, D_FF)\n",
+    "x = torch.randn(BATCH_SIZE, TGT_LEN, D_MODEL)\n",
+    "memory = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
+    "\n",
+    "# Causal mask for decoder self-attention\n",
+    "causal_mask = torch.tril(torch.ones(TGT_LEN, TGT_LEN, dtype=torch.bool)).unsqueeze(0).unsqueeze(0)\n",
+    "out = dec_layer(x, memory, tgt_mask=causal_mask)\n",
+    "assert out.shape == (BATCH_SIZE, TGT_LEN, D_MODEL)\n",
+    "print(f\"✓ DecoderLayer: tgt {x.shape} + memory {memory.shape} -> {out.shape}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9677217c",
+   "metadata": {},
+   "source": [
+    "## 7. Full Transformer (§3)\n",
+    "\n",
+    "> \"Most competitive neural sequence transduction models have an encoder-decoder\n",
+    "> structure. Here, the encoder maps an input sequence of symbol representations\n",
+    "> (x₁, ..., xₙ) to a sequence of continuous representations z = (z₁, ..., zₙ).\n",
+    "> Given z, the decoder then generates an output sequence (y₁, ..., yₘ) of symbols\n",
+    "> one element at a time.\"\n",
+    ">\n",
+    "> — §2, Attention Is All You Need\n",
+    "\n",
+    "> \"In the embedding layers, we multiply those weights by √d_model.\"\n",
+    "> \"In our model, we share the same weight matrix between the two embedding layers\n",
+    "> and the pre-softmax linear transformation.\"\n",
+    ">\n",
+    "> — §3.4, Attention Is All You Need"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1eb4adb5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Transformer(nn.Module):\n",
+    "    \"\"\"§3 — Full Transformer encoder-decoder.\"\"\"\n",
+    "    \n",
+    "    def __init__(self, d_model, n_heads, d_ff, n_layers, vocab_size,\n",
+    "                 max_len=5000, dropout=0.0, norm_eps=1e-6, tie_weights=True):\n",
+    "        super().__init__()\n",
+    "        \n",
+    "        # §3.4 — Embeddings\n",
+    "        self.src_embed = nn.Embedding(vocab_size, d_model)\n",
+    "        self.tgt_embed = nn.Embedding(vocab_size, d_model)\n",
+    "        if tie_weights:  # §3.4 — \"share the same weight matrix\"\n",
+    "            self.tgt_embed.weight = self.src_embed.weight\n",
+    "        \n",
+    "        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)\n",
+    "        self.embed_scale = math.sqrt(d_model)  # §3.4 — multiply by sqrt(d_model)\n",
+    "        \n",
+    "        # §3.1 — Encoder: N identical layers\n",
+    "        self.encoder_layers = nn.ModuleList([\n",
+    "            EncoderLayer(d_model, n_heads, d_ff, dropout, norm_eps)\n",
+    "            for _ in range(n_layers)\n",
+    "        ])\n",
+    "        \n",
+    "        # §3.1 — Decoder: N identical layers\n",
+    "        self.decoder_layers = nn.ModuleList([\n",
+    "            DecoderLayer(d_model, n_heads, d_ff, dropout, norm_eps)\n",
+    "            for _ in range(n_layers)\n",
+    "        ])\n",
+    "        \n",
+    "        # §3.4 — Output projection (tied with embeddings)\n",
+    "        self.output_proj = nn.Linear(d_model, vocab_size, bias=False)\n",
+    "        if tie_weights:\n",
+    "            self.output_proj.weight = self.src_embed.weight\n",
+    "    \n",
+    "    def forward(self, src, tgt, src_mask=None, tgt_mask=None):\n",
+    "        tgt_len = tgt.size(1)\n",
+    "        \n",
+    "        # Auto-generate causal mask if not provided\n",
+    "        if tgt_mask is None:\n",
+    "            tgt_mask = torch.tril(torch.ones(tgt_len, tgt_len, \n",
+    "                         dtype=torch.bool, device=tgt.device)).unsqueeze(0).unsqueeze(0)\n",
+    "        \n",
+    "        # §3.4 — Embed and scale by sqrt(d_model)\n",
+    "        src_emb = self.pos_encoding(self.src_embed(src) * self.embed_scale)\n",
+    "        tgt_emb = self.pos_encoding(self.tgt_embed(tgt) * self.embed_scale)\n",
+    "        \n",
+    "        # Encode\n",
+    "        memory = src_emb\n",
+    "        for layer in self.encoder_layers:\n",
+    "            memory = layer(memory, src_mask)\n",
+    "        \n",
+    "        # Decode\n",
+    "        output = tgt_emb\n",
+    "        for layer in self.decoder_layers:\n",
+    "            output = layer(output, memory, tgt_mask, src_mask)\n",
+    "        \n",
+    "        # §3.4 — Project to vocabulary\n",
+    "        logits = self.output_proj(output)  # (batch, tgt_len, vocab_size)\n",
+    "        return logits"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fc02bf24",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Sanity check: full model forward pass\n",
+    "model = Transformer(\n",
+    "    d_model=D_MODEL, n_heads=N_HEADS, d_ff=D_FF,\n",
+    "    n_layers=N_LAYERS, vocab_size=VOCAB_SIZE,\n",
+    "    dropout=0.0,\n",
+    ")\n",
+    "\n",
+    "src = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SRC_LEN))\n",
+    "tgt = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, TGT_LEN))\n",
+    "logits = model(src, tgt)\n",
+    "\n",
+    "assert logits.shape == (BATCH_SIZE, TGT_LEN, VOCAB_SIZE), \\\n",
+    "    f\"Expected {(BATCH_SIZE, TGT_LEN, VOCAB_SIZE)}, got {logits.shape}\"\n",
+    "print(f\"✓ Full Transformer: src {src.shape} + tgt {tgt.shape} -> logits {logits.shape}\")\n",
+    "\n",
+    "total_params = sum(p.numel() for p in model.parameters())\n",
+    "print(f\"  Parameters: {total_params:,}\")\n",
+    "print(f\"  (Paper base model has ~65M params with d_model=512, vocab=37k)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2e4799ba",
+   "metadata": {},
+   "source": [
+    "## 8. Label-Smoothed Loss (§5.4)\n",
+    "\n",
+    "> \"During training, we employed label smoothing of value ε_ls = 0.1. This hurt\n",
+    "> perplexity, as the model learns to be more unsure, but improved accuracy and\n",
+    "> BLEU score.\"\n",
+    ">\n",
+    "> — §5.4, Attention Is All You Need"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a59bd69",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def label_smoothed_cross_entropy(logits, targets, smoothing=0.1, pad_idx=0):\n",
+    "    \"\"\"§5.4 — Label smoothing with ε_ls = 0.1.\"\"\"\n",
+    "    vocab_size = logits.size(-1)\n",
+    "    logits = logits.contiguous().view(-1, vocab_size)\n",
+    "    targets = targets.contiguous().view(-1)\n",
+    "    \n",
+    "    log_probs = F.log_softmax(logits, dim=-1)\n",
+    "    nll_loss = -log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)\n",
+    "    smooth_loss = -log_probs.mean(dim=-1)\n",
+    "    \n",
+    "    loss = (1.0 - smoothing) * nll_loss + smoothing * smooth_loss\n",
+    "    non_pad = targets != pad_idx\n",
+    "    return (loss * non_pad.float()).sum() / non_pad.float().sum().clamp(min=1.0)\n",
+    "\n",
+    "# Sanity check: loss computes and backprops\n",
+    "targets = torch.randint(1, VOCAB_SIZE, (BATCH_SIZE, TGT_LEN))  # no padding (idx 0)\n",
+    "loss = label_smoothed_cross_entropy(logits, targets, smoothing=0.1)\n",
+    "loss.backward()\n",
+    "print(f\"✓ Label-smoothed loss: {loss.item():.4f}\")\n",
+    "print(f\"  (Compare: -log(1/V) = {-math.log(1/VOCAB_SIZE):.4f} for random predictions)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1174650f",
+   "metadata": {},
+   "source": [
+    "## 9. Gradient Check\n",
+    "\n",
+    "Verify that gradients flow through all parameters of the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da524fdb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Check gradients exist for all parameters\n",
+    "no_grad_params = []\n",
+    "for name, param in model.named_parameters():\n",
+    "    if param.grad is None:\n",
+    "        no_grad_params.append(name)\n",
+    "\n",
+    "if no_grad_params:\n",
+    "    print(f\"⚠ Parameters without gradients: {no_grad_params}\")\n",
+    "else:\n",
+    "    print(f\"✓ All {sum(1 for _ in model.parameters())} parameter tensors have gradients\")\n",
+    "    \n",
+    "# Check gradient norms\n",
+    "total_norm = 0.0\n",
+    "for param in model.parameters():\n",
+    "    if param.grad is not None:\n",
+    "        total_norm += param.grad.data.norm(2).item() ** 2\n",
+    "total_norm = total_norm ** 0.5\n",
+    "print(f\"  Total gradient norm: {total_norm:.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "605c8821",
+   "metadata": {},
+   "source": [
+    "## Common Pitfalls\n",
+    "\n",
+    "1. **Pre-norm vs post-norm**: The paper's Figure 1 shows post-norm (`LayerNorm(x + Sublayer(x))`), and §3.1 says \"followed by layer normalization\" which also implies post-norm. However, many modern reimplementations use pre-norm (`x + Sublayer(LayerNorm(x))`) because it's more stable for deep models. If you're having trouble training, try switching to pre-norm. [§3.1]\n",
+    "\n",
+    "2. **Attention scale factor**: Must divide by √d_k where d_k = d_model/n_heads, NOT √d_model. For the base model: √64 = 8, not √512 ≈ 22.6. This is a 3x difference that significantly affects attention sharpness. [§3.2.1, Eq. 1]\n",
+    "\n",
+    "3. **Positional encoding frequencies**: The div_term uses `exp(arange(0, d_model, 2) * (-log(10000) / d_model))`, stepping by 2. Using `arange(0, d_model)` (step 1) doubles the frequency count and gives wrong wavelengths. [§3.5]\n",
+    "\n",
+    "4. **Label smoothing**: ε_ls = 0.1 means the true class gets probability 1 - ε + ε/V ≈ 0.9 (not exactly 0.9). PyTorch's `CrossEntropyLoss(label_smoothing=0.1)` uses a slightly different formula — for large V this is negligible but worth knowing. [§5.4]\n",
+    "\n",
+    "5. **Learning rate schedule**: The Transformer schedule (Eq. 3) computes the LR from scratch — it's not a multiplier on a base LR. Set the optimizer LR to 1.0 and let the schedule compute the actual LR. Peaks at ~step 4000 for d_model=512. [§5.3, Eq. 3]\n",
+    "\n",
+    "6. **Embedding scaling**: The paper multiplies embedding weights by √d_model (§3.4). This is easy to forget and affects the relative magnitude of embeddings vs positional encodings.\n",
+    "\n",
+    "7. **Weight tying**: The paper shares weights between (a) source embedding, (b) target embedding, and (c) output projection. Missing this changes the parameter count and can affect quality. [§3.4]"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.0.0
+numpy>=1.24.0
+pyyaml>=6.0

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/review.md

@@ -0,0 +1,89 @@
+# Review: Attention Is All You Need — paper2code Output
+
+**Paper:** Attention Is All You Need (Vaswani et al., 2017)
+**Command:** `/paper2code https://arxiv.org/abs/1706.03762`
+**Reviewer:** Human review of skill output
+
+---
+
+## What the skill got right
+
+### Architecture implementation (✓ Correct)
+- Multi-head attention follows §3.2.1 and §3.2.2 precisely
+- Scaled dot-product attention implements Eq. 1 correctly: softmax(QK^T/√d_k)V
+- Scale factor uses d_k (= d_model/n_heads = 64), not d_model — this is correct and a common mistake to catch
+- Positional encoding matches §3.5 exactly — sinusoidal with correct div_term formula
+- Encoder/decoder stacks are N=6 layers as specified in Table 3
+- Post-norm (LayerNorm after residual addition) matches Figure 1
+
+### Citation anchoring (✓ Thorough)
+- Every class and function references the correct paper section
+- Equation numbers match the paper's numbering
+- Shape comments are present on every tensor operation
+- Variable names match paper notation (Q, K, V, d_k, d_model)
+
+### Ambiguity flagging (✓ Honest)
+- LayerNorm epsilon flagged as [UNSPECIFIED] — correct, paper doesn't mention it
+- Weight initialization flagged as [UNSPECIFIED] — correct
+- Bias in linear layers flagged as [UNSPECIFIED] — correct
+- Max sequence length for PE flagged as [UNSPECIFIED] — correct
+
+### Config file (✓ Complete)
+- Every parameter has a paper citation or [UNSPECIFIED] flag
+- β₁, β₂, ε for Adam correctly sourced from §5.3
+- Label smoothing ε=0.1 from §5.4
+- Warmup steps 4000 from §5.3
+
+---
+
+## What it correctly flagged as unspecified
+
+1. **LayerNorm epsilon** — the paper never states this. Using 1e-6 is a reasonable choice. ✓
+2. **Bias terms in linear projections** — never discussed in the paper. ✓
+3. **Weight initialization** — the paper does not describe initialization. Xavier uniform is a reasonable default. ✓
+4. **Max sequence length for positional encoding** — paper doesn't specify an upper bound beyond the training data. ✓
+5. **Gradient clipping** — paper doesn't mention it. ✓
+6. **Dropout on attention weights specifically** — paper says P_drop=0.1 dropout generally but doesn't enumerate every placement. ✓
+
+---
+
+## Potential issues and discussion points
+
+### 1. Pre-norm vs Post-norm ambiguity
+The implementation uses post-norm (LayerNorm after residual addition), matching Figure 1. However, some researchers argue the paper describes pre-norm in the text. The implementation correctly notes this in REPRODUCTION_NOTES.md but could be more explicit about the ambiguity in the code comments.
+
+**Verdict:** Acceptable. The choice is reasonable and the ambiguity is documented.
+
+### 2. Dropout placement
+The paper says "Residual Dropout" in §5.4 and describes it as being applied "to the output of each sub-layer, before it is added to the sub-layer input and normalized." The implementation applies dropout after attention and after FFN, plus on embeddings. The exact dropout placement across attention weights, attention output, FFN internal, FFN output, and embeddings is somewhat ambiguous.
+
+**Verdict:** The implementation is reasonable but could have flagged the attention weight dropout more specifically as a partial specification.
+
+### 3. Weight tying
+The implementation correctly implements three-way weight tying (source embedding, target embedding, output projection) as described in §3.4. This is a detail many implementations miss — good catch.
+
+### 4. Embedding scaling
+The √d_model scaling on embeddings (§3.4) is correctly implemented. This is another commonly missed detail.
+
+### 5. Training scope
+Correctly identified that this is an architecture paper (type a) and appropriately limited the training code to a minimal example rather than a full pipeline. Good scope decision.
+
+---
+
+## What could be improved
+
+1. **The walkthrough notebook** — would benefit from more explicit sanity checks on intermediate tensor shapes, especially showing that attention weights sum to 1 after softmax.
+
+2. **BLEU implementation note** — could mention that the original paper likely used multi-bleu.perl, which gives slightly different numbers than sacrebleu.
+
+3. **Beam search** — the paper's reported numbers use beam search with beam_size=4 and length_penalty=0.6 (§6.1). A note about this in the training script would be helpful for anyone trying to reproduce Table 2.
+
+---
+
+## Overall assessment
+
+**Quality: High.** The implementation is citation-anchored, honest about uncertainty, and correctly scoped. A researcher could use this as a starting point for reproduction with high confidence that the architecture matches the paper. The REPRODUCTION_NOTES.md is genuinely useful — it saves someone the work of auditing the paper for unspecified details.
+
+**Coverage: Core contribution fully implemented.** All architectural components from §3 are present and correct. Training infrastructure is appropriately minimal (this is an architecture paper, not a training method paper).
+
+**Honesty: Good.** Unspecified choices are flagged consistently. No silent assumptions detected. The review found no cases where the implementation confidently invented a detail that the paper doesn't specify.

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/src/data.py

@@ -0,0 +1,156 @@
+"""
+Attention Is All You Need — Dataset and Data Loading
+
+Paper: https://arxiv.org/abs/1706.03762
+Implements: Data loading skeleton for WMT translation tasks.
+
+Section references:
+  §5.1 — "We trained on the standard WMT 2014 English-German dataset
+          consisting of about 4.5 million sentence pairs."
+  §5.1 — "Sentences were encoded using byte-pair encoding, which has a shared
+          source-target vocabulary of about 37000 tokens."
+
+NOTE: This file provides the Dataset class skeleton. You must:
+  1. Download WMT 2014 EN-DE from http://www.statmt.org/wmt14/translation-task.html
+  2. Apply BPE tokenization (e.g., sentencepiece with 37000 merge operations)
+  3. Set the data_dir in configs/base.yaml
+"""
+
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+
+class WMTTranslationDataset(Dataset):
+    """§5.1 — WMT 2014 English-German translation dataset.
+
+    "We trained on the standard WMT 2014 English-German dataset consisting of
+     about 4.5 million sentence pairs. Sentences were encoded using byte-pair
+     encoding, which has a shared source-target vocabulary of about 37000 tokens."
+
+    Expected data format:
+        {data_dir}/
+        ├── train.src     # Source sentences, one per line, BPE-tokenized
+        ├── train.tgt     # Target sentences, one per line, BPE-tokenized
+        ├── valid.src
+        ├── valid.tgt
+        ├── test.src
+        ├── test.tgt
+        └── vocab.txt     # Shared BPE vocabulary, one token per line
+
+    How to obtain:
+        1. Download WMT 2014 EN-DE data from http://www.statmt.org/wmt14/
+        2. Apply BPE tokenization with ~37000 merge operations
+           (use sentencepiece: spm_train --input=data.txt --model_prefix=wmt
+            --vocab_size=37000 --model_type=bpe)
+        3. Encode source and target files with the learned BPE model
+
+    Preprocessing (§5.1):
+        - Byte-pair encoding with shared source-target vocabulary
+        - ~37000 BPE tokens for EN-DE
+        - ~32000 BPE tokens for EN-FR (word-piece model)
+    """
+
+    def __init__(
+        self,
+        data_dir: str,
+        split: str = "train",
+        max_seq_len: int = 512,
+        pad_idx: int = 0,
+        bos_idx: int = 1,
+        eos_idx: int = 2,
+    ):
+        """
+        Args:
+            data_dir: path to the preprocessed data directory
+            split: one of "train", "valid", "test"
+            max_seq_len: maximum sequence length (truncate longer sequences)
+            pad_idx: padding token index
+            bos_idx: beginning-of-sentence token index
+            eos_idx: end-of-sentence token index
+        """
+        self.data_dir = Path(data_dir)
+        self.split = split
+        self.max_seq_len = max_seq_len
+        self.pad_idx = pad_idx
+        self.bos_idx = bos_idx
+        self.eos_idx = eos_idx
+
+        # TODO: Load tokenized data
+        # self.src_data, self.tgt_data = self._load_data()
+        self.src_data: List[List[int]] = []
+        self.tgt_data: List[List[int]] = []
+
+    def _load_data(self) -> Tuple[List[List[int]], List[List[int]]]:
+        """Load BPE-tokenized source and target sentences.
+
+        TODO: Implement based on your tokenization format.
+        Expected: each line is space-separated token IDs.
+        """
+        src_path = self.data_dir / f"{self.split}.src"
+        tgt_path = self.data_dir / f"{self.split}.tgt"
+
+        if not src_path.exists() or not tgt_path.exists():
+            raise FileNotFoundError(
+                f"Data files not found: {src_path}, {tgt_path}\n"
+                f"Please download and preprocess the WMT 2014 dataset.\n"
+                f"See docstring for instructions."
+            )
+
+        # TODO: Parse the tokenized files
+        raise NotImplementedError(
+            "Implement _load_data() to read your BPE-tokenized data files."
+        )
+
+    def __len__(self) -> int:
+        return len(self.src_data)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """Load and preprocess a single sentence pair.
+
+        Returns:
+            dict with keys:
+                "src": source token IDs — shape: (src_len,)
+                "tgt": target token IDs (with BOS prepended) — shape: (tgt_len,)
+                "tgt_labels": target labels (with EOS appended) — shape: (tgt_len,)
+        """
+        src_tokens = self.src_data[idx][:self.max_seq_len]
+        tgt_tokens = self.tgt_data[idx][:self.max_seq_len - 1]  # leave room for BOS/EOS
+
+        # Decoder input: [BOS] + target tokens
+        tgt_input = [self.bos_idx] + tgt_tokens
+        # Decoder labels: target tokens + [EOS]
+        tgt_labels = tgt_tokens + [self.eos_idx]
+
+        return {
+            "src": torch.tensor(src_tokens, dtype=torch.long),
+            "tgt": torch.tensor(tgt_input, dtype=torch.long),
+            "tgt_labels": torch.tensor(tgt_labels, dtype=torch.long),
+        }
+
+
+def collate_fn(batch: List[Dict[str, torch.Tensor]], pad_idx: int = 0) -> Dict[str, torch.Tensor]:
+    """Pad sequences to the same length within a batch.
+
+    §5.1 — "Each training batch contained a set of sentence pairs containing
+    approximately 25000 source tokens and 25000 target tokens."
+
+    NOTE: The paper batches by approximate token count, not by number of sentences.
+    This collate function pads by sequence length for simplicity.
+    For exact reproduction, implement token-count-based batching.
+    """
+    src_tensors = [item["src"] for item in batch]
+    tgt_tensors = [item["tgt"] for item in batch]
+    label_tensors = [item["tgt_labels"] for item in batch]
+
+    src_padded = torch.nn.utils.rnn.pad_sequence(src_tensors, batch_first=True, padding_value=pad_idx)
+    tgt_padded = torch.nn.utils.rnn.pad_sequence(tgt_tensors, batch_first=True, padding_value=pad_idx)
+    labels_padded = torch.nn.utils.rnn.pad_sequence(label_tensors, batch_first=True, padding_value=pad_idx)
+
+    return {
+        "src": src_padded,
+        "tgt": tgt_padded,
+        "tgt_labels": labels_padded,
+    }

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/src/evaluate.py

@@ -0,0 +1,73 @@
+"""
+Attention Is All You Need — Evaluation Metrics
+
+Paper: https://arxiv.org/abs/1706.03762
+Implements: BLEU score computation for translation evaluation.
+
+Section references:
+  §6 — "Results" — BLEU scores reported in Table 2
+
+NOTE: For reproducible BLEU, use sacrebleu (pip install sacrebleu).
+      Different BLEU implementations give different numbers (sometimes 1-2 points)
+      due to tokenization and smoothing differences.
+"""
+
+from typing import Dict, List
+
+
+def compute_bleu(
+    hypotheses: List[str],
+    references: List[List[str]],
+) -> float:
+    """Compute BLEU score using sacrebleu for reproducibility.
+
+    §6, Table 2 — "BLEU" score on WMT 2014 EN-DE and EN-FR.
+
+    The paper reports BLEU scores of:
+      - 27.3 BLEU on EN-DE (base model)
+      - 28.4 BLEU on EN-DE (big model)
+      - 38.1 BLEU on EN-FR (big model)
+
+    These scores use beam search with beam size 4 and length penalty α=0.6 (§6.1).
+
+    Args:
+        hypotheses: list of predicted translations (detokenized strings)
+        references: list of lists of reference translations
+
+    Returns:
+        BLEU score as a float
+
+    NOTE: [UNSPECIFIED] The paper does not specify which BLEU implementation
+    was used. We use sacrebleu (Post, 2018) for reproducibility, which is the
+    current standard. Original 2017 results may have used multi-bleu.perl.
+    """
+    try:
+        import sacrebleu
+        # sacrebleu expects: hypotheses as list[str], references as list[list[str]]
+        bleu = sacrebleu.corpus_bleu(hypotheses, references)
+        return bleu.score
+    except ImportError:
+        raise ImportError(
+            "sacrebleu is required for BLEU computation. "
+            "Install with: pip install sacrebleu"
+        )
+
+
+def compute_all_metrics(
+    hypotheses: List[str],
+    references: List[List[str]],
+) -> Dict[str, float]:
+    """Compute all evaluation metrics reported in the paper.
+
+    §6, Table 2 — "We report BLEU scores"
+
+    Args:
+        hypotheses: list of predicted translations
+        references: list of lists of reference translations
+
+    Returns:
+        Dict mapping metric name to value
+    """
+    return {
+        "BLEU": compute_bleu(hypotheses, references),
+    }

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/src/loss.py

@@ -0,0 +1,78 @@
+"""
+Attention Is All You Need — Loss Functions
+
+Paper: https://arxiv.org/abs/1706.03762
+Implements: Label-smoothed cross-entropy loss
+
+Section references:
+  §5.4 — "During training, we employed label smoothing of value ε_ls = 0.1"
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LabelSmoothedCrossEntropy(nn.Module):
+    """§5.4 — Label smoothing with KL divergence.
+
+    "During training, we employed label smoothing of value ε_ls = 0.1.
+     This hurt perplexity, as the model learns to be more unsure,
+     but improved accuracy and BLEU score."
+
+    Uses KL divergence formulation following Szegedy et al. (2016).
+    Target distribution: (1 - ε) on correct class, ε / (V-1) on other classes.
+
+    NOTE: This differs slightly from PyTorch's built-in label_smoothing parameter
+    in CrossEntropyLoss. PyTorch distributes ε uniformly as ε/V to all classes
+    (including correct), while the paper's formulation distributes to incorrect
+    classes only. For large vocabularies (V >> 1), the difference is negligible.
+    We implement the paper's stated formulation.
+    """
+
+    def __init__(self, smoothing: float = 0.1, pad_idx: int = 0):
+        """
+        Args:
+            smoothing: §5.4 — ε_ls = 0.1
+            pad_idx: index of padding token (loss ignored for these positions)
+        """
+        super().__init__()
+        self.smoothing = smoothing        # §5.4 — "ε_ls = 0.1"
+        self.confidence = 1.0 - smoothing  # probability mass on correct class
+        self.pad_idx = pad_idx
+
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """Compute label-smoothed cross-entropy loss.
+
+        Args:
+            logits: (batch, seq_len, vocab_size) — model output logits (NOT probabilities)
+            targets: (batch, seq_len) — target token indices
+
+        Returns:
+            Scalar loss value
+        """
+        vocab_size = logits.size(-1)
+
+        # Reshape for computation
+        logits = logits.contiguous().view(-1, vocab_size)   # (batch * seq_len, vocab_size)
+        targets = targets.contiguous().view(-1)              # (batch * seq_len,)
+
+        # Compute log probabilities (numerically stable via log_softmax)
+        log_probs = F.log_softmax(logits, dim=-1)  # (batch * seq_len, vocab_size)
+
+        # NLL loss on the true class
+        nll_loss = -log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)
+        # (batch * seq_len,)
+
+        # Smooth loss: uniform distribution over all classes
+        smooth_loss = -log_probs.mean(dim=-1)  # (batch * seq_len,)
+
+        # §5.4 — Combined loss with label smoothing
+        loss = self.confidence * nll_loss + self.smoothing * smooth_loss
+
+        # Mask out padding positions
+        non_pad_mask = targets != self.pad_idx
+        loss = loss * non_pad_mask.float()
+
+        # Average over non-padding tokens
+        return loss.sum() / non_pad_mask.float().sum().clamp(min=1.0)

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/src/model.py

@@ -0,0 +1,569 @@
+"""
+Attention Is All You Need — Model Architecture
+
+Paper: https://arxiv.org/abs/1706.03762
+Authors: Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, Kaiser, Polosukhin (2017)
+
+Implements: The Transformer model — an encoder-decoder architecture based entirely
+on multi-head self-attention, with no recurrence or convolution.
+
+Section references:
+  §3.1 — Encoder and Decoder Stacks
+  §3.2 — Attention (Scaled Dot-Product and Multi-Head)
+  §3.3 — Position-wise Feed-Forward Networks
+  §3.4 — Embeddings and Softmax
+  §3.5 — Positional Encoding
+
+Usage:
+    from src.model import Transformer, TransformerConfig
+
+    config = TransformerConfig()
+    model = Transformer(config)
+    output = model(src_tokens, tgt_tokens)
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+@dataclass
+class TransformerConfig:
+    """All model hyperparameters.
+
+    Values from Vaswani et al. 2017, Table 3 ("base model") unless marked [UNSPECIFIED].
+    """
+    # Architecture — §3, Table 3
+    d_model: int = 512          # §3.1, Table 3 — "d_model = 512"
+    n_heads: int = 8            # §3.2.2, Table 3 — "h = 8"
+    d_ff: int = 2048            # §3.3, Table 3 — "d_ff = 2048"
+    n_encoder_layers: int = 6   # §3.1, Table 3 — "N = 6"
+    n_decoder_layers: int = 6   # §3.1, Table 3 — "N = 6"
+    dropout: float = 0.1        # §5.4, Table 3 — "P_drop = 0.1"
+    vocab_size: int = 37000     # §5.1 — "~37000 tokens" (EN-DE BPE)
+    max_seq_len: int = 5000     # [UNSPECIFIED] — max length for positional encoding
+    norm_eps: float = 1e-6      # [UNSPECIFIED] — LayerNorm epsilon not stated. Alternatives: 1e-5 (PyTorch default)
+    tie_weights: bool = True    # §3.4 — "we share the same weight matrix"
+
+
+# ---------------------------------------------------------------------------
+# §3.2.1 — Scaled Dot-Product Attention
+# ---------------------------------------------------------------------------
+
+def scaled_dot_product_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    dropout: Optional[nn.Dropout] = None,
+) -> torch.Tensor:
+    """§3.2.1, Eq. 1 — Attention(Q, K, V) = softmax(QK^T / √d_k) V
+
+    "We compute the attention function on a set of queries simultaneously,
+     packed together into a matrix Q."
+
+    Args:
+        query: (batch, n_heads, seq_q, d_k)
+        key:   (batch, n_heads, seq_k, d_k)
+        value: (batch, n_heads, seq_k, d_v)
+        mask:  (batch, 1, seq_q, seq_k) or (1, 1, seq_q, seq_k) — True = attend, False = mask
+        dropout: optional dropout on attention weights
+
+    Returns:
+        (batch, n_heads, seq_q, d_v) — attention output
+    """
+    d_k = query.size(-1)
+
+    # §3.2.1 — "We compute the dot products of the query with all keys,
+    # divide each by √d_k, and apply a softmax"
+    # (batch, n_heads, seq_q, d_k) @ (batch, n_heads, d_k, seq_k) -> (batch, n_heads, seq_q, seq_k)
+    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
+
+    if mask is not None:
+        # §3.2.3 — "We implement this inside of scaled dot-product attention by masking out
+        # (setting to −∞) all values in the input of the softmax which correspond to
+        # illegal connections"
+        scores = scores.masked_fill(mask == 0, float("-inf"))
+
+    # (batch, n_heads, seq_q, seq_k)
+    attn_weights = F.softmax(scores, dim=-1)
+
+    if dropout is not None:
+        attn_weights = dropout(attn_weights)
+
+    # (batch, n_heads, seq_q, seq_k) @ (batch, n_heads, seq_k, d_v) -> (batch, n_heads, seq_q, d_v)
+    return torch.matmul(attn_weights, value)
+
+
+# ---------------------------------------------------------------------------
+# §3.2.2 — Multi-Head Attention
+# ---------------------------------------------------------------------------
+
+class MultiHeadAttention(nn.Module):
+    """§3.2.2 — Multi-Head Attention.
+
+    "Multi-head attention allows the model to jointly attend to information
+     from different representation subspaces at different positions."
+
+    MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O
+    where head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)
+    """
+
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        assert config.d_model % config.n_heads == 0, \
+            f"d_model ({config.d_model}) must be divisible by n_heads ({config.n_heads})"
+
+        self.d_model = config.d_model
+        self.n_heads = config.n_heads
+        self.d_k = config.d_model // config.n_heads  # §3.2.2 — "d_k = d_v = d_model/h = 64"
+
+        # §3.2.2 — "We employ h = 8 parallel attention layers, or heads"
+        # "We use d_k = d_v = d_model/h = 64"
+        # [UNSPECIFIED] Paper does not state whether projection layers have bias.
+        # Using: bias=True (common default)
+        # Alternatives: bias=False (used in some reimplementations like LLaMA)
+        self.W_q = nn.Linear(config.d_model, config.d_model)
+        self.W_k = nn.Linear(config.d_model, config.d_model)
+        self.W_v = nn.Linear(config.d_model, config.d_model)
+        self.W_o = nn.Linear(config.d_model, config.d_model)
+
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            query: (batch, seq_q, d_model)
+            key:   (batch, seq_k, d_model)
+            value: (batch, seq_k, d_model)
+            mask:  (batch, 1, seq_q, seq_k) or broadcastable shape
+
+        Returns:
+            (batch, seq_q, d_model)
+        """
+        batch_size = query.size(0)
+
+        # §3.2.2 — "linearly project the queries, keys and values h times with different,
+        # learned linear projections to d_k, d_k, and d_v dimensions"
+        # (batch, seq, d_model) -> (batch, n_heads, seq, d_k)
+        q = self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+        k = self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+        v = self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+
+        # §3.2.1 — apply scaled dot-product attention
+        attn_out = scaled_dot_product_attention(q, k, v, mask=mask, dropout=self.dropout)
+        # (batch, n_heads, seq_q, d_k)
+
+        # §3.2.2 — "concatenated and once again projected"
+        # (batch, n_heads, seq_q, d_k) -> (batch, seq_q, d_model)
+        attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
+
+        # Final linear projection W^O
+        return self.W_o(attn_out)  # (batch, seq_q, d_model)
+
+
+# ---------------------------------------------------------------------------
+# §3.3 — Position-wise Feed-Forward Networks
+# ---------------------------------------------------------------------------
+
+class FeedForward(nn.Module):
+    """§3.3 — Position-wise Feed-Forward Networks.
+
+    "FFN(x) = max(0, xW_1 + b_1)W_2 + b_2" (Eq. 2)
+
+    "The dimensionality of input and output is d_model = 512,
+     and the inner-layer has dimensionality d_ff = 2048."
+    """
+
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        # §3.3, Eq. 2 — two linear transformations with ReLU in between
+        self.linear1 = nn.Linear(config.d_model, config.d_ff)
+        self.linear2 = nn.Linear(config.d_ff, config.d_model)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (batch, seq_len, d_model)
+        Returns:
+            (batch, seq_len, d_model)
+        """
+        # §3.3, Eq. 2 — FFN(x) = max(0, xW_1 + b_1)W_2 + b_2
+        x = self.linear1(x)          # (batch, seq_len, d_model) -> (batch, seq_len, d_ff)
+        x = F.relu(x)                # §3.3, Eq. 2 — ReLU activation
+        x = self.dropout(x)          # §5.4 — "applied... to the output of each sub-layer"
+        x = self.linear2(x)          # (batch, seq_len, d_ff) -> (batch, seq_len, d_model)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# §3.5 — Positional Encoding
+# ---------------------------------------------------------------------------
+
+class PositionalEncoding(nn.Module):
+    """§3.5 — Positional Encoding (sinusoidal).
+
+    "We use sine and cosine functions of different frequencies:"
+      PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
+      PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
+
+    "We chose this function because we hypothesized it would allow the model
+     to easily learn to attend by relative positions."
+    """
+
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        self.dropout = nn.Dropout(config.dropout)
+
+        # §3.5 — precompute positional encodings
+        pe = torch.zeros(config.max_seq_len, config.d_model)               # (max_len, d_model)
+        position = torch.arange(0, config.max_seq_len).unsqueeze(1).float()  # (max_len, 1)
+        div_term = torch.exp(
+            torch.arange(0, config.d_model, 2).float()
+            * (-math.log(10000.0) / config.d_model)
+        )  # (d_model/2,)
+
+        pe[:, 0::2] = torch.sin(position * div_term)  # even indices: sin
+        pe[:, 1::2] = torch.cos(position * div_term)  # odd indices: cos
+        pe = pe.unsqueeze(0)  # (1, max_len, d_model) — for batch broadcasting
+        self.register_buffer("pe", pe)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (batch, seq_len, d_model) — embedding output (already scaled by √d_model)
+        Returns:
+            (batch, seq_len, d_model) — with positional encoding added
+        """
+        # §3.5 — "We apply dropout to the sums of the embeddings and the positional encodings"
+        x = x + self.pe[:, :x.size(1)]
+        return self.dropout(x)
+
+
+# ---------------------------------------------------------------------------
+# §3.1 — Encoder Layer
+# ---------------------------------------------------------------------------
+
+class EncoderLayer(nn.Module):
+    """§3.1 — Single encoder layer.
+
+    "Each layer has two sub-layers. The first is a multi-head self-attention mechanism,
+     and the second is a simple, position-wise fully connected feed-forward network."
+
+    "We employ a residual connection around each of the two sub-layers,
+     followed by layer normalization."
+
+    LayerNorm(x + Sublayer(x))  — post-norm formulation
+    """
+
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        # §3.1 — two sub-layers with residual connections and layer norm
+        self.self_attn = MultiHeadAttention(config)
+        self.feed_forward = FeedForward(config)
+
+        # §3.1 — "layer normalization" after each residual connection
+        self.norm1 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
+        self.norm2 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
+
+        # §5.4 — "Residual Dropout: We apply dropout to the output of each sub-layer,
+        # before it is added to the sub-layer input and normalized"
+        self.dropout1 = nn.Dropout(config.dropout)
+        self.dropout2 = nn.Dropout(config.dropout)
+
+    def forward(
+        self, x: torch.Tensor, src_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: (batch, seq_len, d_model)
+            src_mask: (batch, 1, 1, seq_len) — padding mask
+        Returns:
+            (batch, seq_len, d_model)
+        """
+        # §3.1 — Sub-layer 1: Multi-head self-attention
+        # Post-norm: LayerNorm(x + Dropout(Sublayer(x)))
+        attn_out = self.self_attn(x, x, x, mask=src_mask)  # (batch, seq_len, d_model)
+        x = self.norm1(x + self.dropout1(attn_out))          # (batch, seq_len, d_model)
+
+        # §3.1 — Sub-layer 2: Feed-forward network
+        ff_out = self.feed_forward(x)                         # (batch, seq_len, d_model)
+        x = self.norm2(x + self.dropout2(ff_out))             # (batch, seq_len, d_model)
+
+        return x
+
+
+# ---------------------------------------------------------------------------
+# §3.1 — Decoder Layer
+# ---------------------------------------------------------------------------
+
+class DecoderLayer(nn.Module):
+    """§3.1 — Single decoder layer.
+
+    "In addition to the two sub-layers in each encoder layer, the decoder inserts
+     a third sub-layer, which performs multi-head attention over the output of the
+     encoder stack."
+
+    Three sub-layers:
+      1. Masked multi-head self-attention (causal)
+      2. Multi-head cross-attention (over encoder output)
+      3. Position-wise feed-forward network
+    """
+
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        # §3.1 — three sub-layers
+        self.self_attn = MultiHeadAttention(config)
+        self.cross_attn = MultiHeadAttention(config)
+        self.feed_forward = FeedForward(config)
+
+        self.norm1 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
+        self.norm2 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
+        self.norm3 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
+
+        self.dropout1 = nn.Dropout(config.dropout)
+        self.dropout2 = nn.Dropout(config.dropout)
+        self.dropout3 = nn.Dropout(config.dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        memory: torch.Tensor,
+        tgt_mask: Optional[torch.Tensor] = None,
+        memory_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: (batch, tgt_len, d_model) — decoder input
+            memory: (batch, src_len, d_model) — encoder output
+            tgt_mask: (1, 1, tgt_len, tgt_len) — causal mask
+            memory_mask: (batch, 1, 1, src_len) — source padding mask
+        Returns:
+            (batch, tgt_len, d_model)
+        """
+        # §3.1 — Sub-layer 1: Masked self-attention
+        # §3.2.3 — "We need to prevent leftward information flow in the decoder
+        # to preserve the auto-regressive property"
+        self_attn_out = self.self_attn(x, x, x, mask=tgt_mask)
+        x = self.norm1(x + self.dropout1(self_attn_out))  # (batch, tgt_len, d_model)
+
+        # §3.1 — Sub-layer 2: Cross-attention over encoder output
+        # "queries come from the previous decoder layer, and the memory keys and values
+        # come from the output of the encoder"
+        cross_attn_out = self.cross_attn(x, memory, memory, mask=memory_mask)
+        x = self.norm2(x + self.dropout2(cross_attn_out))  # (batch, tgt_len, d_model)
+
+        # §3.1 — Sub-layer 3: Feed-forward
+        ff_out = self.feed_forward(x)
+        x = self.norm3(x + self.dropout3(ff_out))  # (batch, tgt_len, d_model)
+
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Encoder and Decoder Stacks
+# ---------------------------------------------------------------------------
+
+class Encoder(nn.Module):
+    """§3.1 — Encoder: "The encoder is composed of a stack of N = 6 identical layers." """
+
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            EncoderLayer(config) for _ in range(config.n_encoder_layers)
+        ])
+
+    def forward(
+        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: (batch, src_len, d_model) — embedded + positionally-encoded source
+            mask: (batch, 1, 1, src_len) — source padding mask
+        Returns:
+            (batch, src_len, d_model) — encoder output (memory)
+        """
+        for layer in self.layers:
+            x = layer(x, src_mask=mask)
+        return x
+
+
+class Decoder(nn.Module):
+    """§3.1 — Decoder: "The decoder is also composed of a stack of N = 6 identical layers." """
+
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            DecoderLayer(config) for _ in range(config.n_decoder_layers)
+        ])
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        memory: torch.Tensor,
+        tgt_mask: Optional[torch.Tensor] = None,
+        memory_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: (batch, tgt_len, d_model) — embedded + positionally-encoded target
+            memory: (batch, src_len, d_model) — encoder output
+            tgt_mask: (1, 1, tgt_len, tgt_len) — causal mask
+            memory_mask: (batch, 1, 1, src_len) — source padding mask
+        Returns:
+            (batch, tgt_len, d_model)
+        """
+        for layer in self.layers:
+            x = layer(x, memory, tgt_mask=tgt_mask, memory_mask=memory_mask)
+        return x
+
+
+# ---------------------------------------------------------------------------
+# §3 — Full Transformer Model
+# ---------------------------------------------------------------------------
+
+class Transformer(nn.Module):
+    """§3 — The Transformer (Encoder-Decoder).
+
+    "Most competitive neural sequence transduction models have an encoder-decoder structure.
+     Here, the encoder maps an input sequence of symbol representations (x_1, ..., x_n) to
+     a sequence of continuous representations z = (z_1, ..., z_n). Given z, the decoder then
+     generates an output sequence (y_1, ..., y_m) of symbols one element at a time."
+
+    Composed of:
+      - Token embedding with √d_model scaling (§3.4)
+      - Positional encoding (§3.5)
+      - Encoder stack — N=6 layers (§3.1)
+      - Decoder stack — N=6 layers (§3.1)
+      - Linear output projection with weight tying (§3.4)
+    """
+
+    def __init__(self, config: TransformerConfig):
+        super().__init__()
+        self.config = config
+
+        # §3.4 — "Similarly to other sequence transduction models, we use learned embeddings
+        # to convert the input tokens and output tokens to vectors of dimension d_model"
+        self.src_embedding = nn.Embedding(config.vocab_size, config.d_model)
+        self.tgt_embedding = nn.Embedding(config.vocab_size, config.d_model)
+
+        # §3.4 — "In our model, we share the same weight matrix between the two
+        # embedding layers and the pre-softmax linear transformation"
+        if config.tie_weights:
+            self.tgt_embedding.weight = self.src_embedding.weight
+
+        # §3.5 — Positional encoding (shared between encoder and decoder)
+        self.pos_encoding = PositionalEncoding(config)
+
+        # §3.1 — Encoder and decoder stacks
+        self.encoder = Encoder(config)
+        self.decoder = Decoder(config)
+
+        # §3.4 — Output linear projection
+        # "we share the same weight matrix between the two embedding layers
+        # and the pre-softmax linear transformation"
+        self.output_projection = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        if config.tie_weights:
+            self.output_projection.weight = self.src_embedding.weight
+
+        # §3.4 — embedding scale factor: √d_model
+        self.embed_scale = math.sqrt(config.d_model)
+
+        # [UNSPECIFIED] Paper does not describe weight initialization.
+        # Using: Xavier uniform for all linear layers (common for transformers)
+        # Alternatives: PyTorch defaults, normal(0, 0.02)
+        self._init_weights()
+
+    def _init_weights(self):
+        """[UNSPECIFIED] — Weight initialization not described in paper.
+        Using Xavier uniform, a common choice for transformer models.
+        """
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        tgt: torch.Tensor,
+        src_mask: Optional[torch.Tensor] = None,
+        tgt_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Full encoder-decoder forward pass.
+
+        Args:
+            src: (batch, src_len) — source token IDs
+            tgt: (batch, tgt_len) — target token IDs (shifted right for teacher forcing)
+            src_mask: (batch, 1, 1, src_len) — source padding mask. Auto-generated if None.
+            tgt_mask: (1, 1, tgt_len, tgt_len) — causal mask. Auto-generated if None.
+
+        Returns:
+            (batch, tgt_len, vocab_size) — logits over vocabulary
+        """
+        # Auto-generate masks if not provided
+        if src_mask is None:
+            src_mask = (src != 0).unsqueeze(1).unsqueeze(2)  # (batch, 1, 1, src_len)
+        if tgt_mask is None:
+            tgt_mask = self._make_causal_mask(tgt.size(1), tgt.device)  # (1, 1, tgt_len, tgt_len)
+
+        # §3.4 — "In the embedding layers, we multiply those weights by √d_model"
+        src_emb = self.src_embedding(src) * self.embed_scale  # (batch, src_len, d_model)
+        tgt_emb = self.tgt_embedding(tgt) * self.embed_scale  # (batch, tgt_len, d_model)
+
+        # §3.5 — Add positional encoding
+        src_emb = self.pos_encoding(src_emb)  # (batch, src_len, d_model)
+        tgt_emb = self.pos_encoding(tgt_emb)  # (batch, tgt_len, d_model)
+
+        # §3.1 — Encode
+        memory = self.encoder(src_emb, mask=src_mask)  # (batch, src_len, d_model)
+
+        # §3.1 — Decode
+        dec_out = self.decoder(
+            tgt_emb, memory,
+            tgt_mask=tgt_mask,
+            memory_mask=src_mask,
+        )  # (batch, tgt_len, d_model)
+
+        # §3.4 — Linear projection to vocabulary
+        logits = self.output_projection(dec_out)  # (batch, tgt_len, vocab_size)
+        return logits
+
+    @staticmethod
+    def _make_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
+        """§3.2.3 — Create causal (autoregressive) mask for decoder self-attention.
+
+        "We also modify the self-attention sub-layer in the decoder stack to prevent
+         positions from attending to subsequent positions."
+
+        Returns:
+            (1, 1, seq_len, seq_len) boolean tensor. True = attend, False = mask.
+        """
+        mask = torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool))
+        return mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, seq_len)
+
+    def __repr__(self) -> str:
+        total_params = sum(p.numel() for p in self.parameters())
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        return (
+            f"Transformer(\n"
+            f"  config=TransformerConfig(d_model={self.config.d_model}, n_heads={self.config.n_heads}, "
+            f"d_ff={self.config.d_ff}, n_layers={self.config.n_encoder_layers}),\n"
+            f"  total_params={total_params:,},\n"
+            f"  trainable_params={trainable_params:,}\n"
+            f")"
+        )

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/src/train.py

@@ -0,0 +1,97 @@
+"""
+Attention Is All You Need — Minimal Training Example
+
+Paper: https://arxiv.org/abs/1706.03762
+Implements: Minimal training loop demonstrating optimizer setup and LR schedule.
+
+This is NOT a full training pipeline. The paper's contribution is primarily
+architectural (§3). This script shows how to set up training with the
+paper's specified optimizer and schedule (§5.3) for reference.
+
+For actual training, you need:
+  - WMT 2014 data with BPE tokenization (§5.1)
+  - 8 P100 GPUs for the reported training speed (§5.2)
+  - Token-count batching (~25k source + 25k target tokens per batch) (§5.1)
+"""
+
+import math
+
+import torch
+import yaml
+
+from src.model import Transformer, TransformerConfig
+from src.loss import LabelSmoothedCrossEntropy
+from src.utils import get_transformer_lr_schedule
+
+
+def train(config_path: str = "configs/base.yaml"):
+    """Minimal training example with paper-specified optimizer and schedule."""
+    # Load config
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Build model — §3, Table 3
+    model_config = TransformerConfig(
+        d_model=config["model"]["d_model"],
+        n_heads=config["model"]["n_heads"],
+        d_ff=config["model"]["d_ff"],
+        n_encoder_layers=config["model"]["n_encoder_layers"],
+        n_decoder_layers=config["model"]["n_decoder_layers"],
+        dropout=config["model"]["dropout"],
+        vocab_size=config["model"]["vocab_size"],
+        norm_eps=config["model"]["norm_eps"],
+        tie_weights=config["model"]["tie_weights"],
+    )
+    model = Transformer(model_config).to(device)
+    print(model)
+
+    # §5.3 — "We used the Adam optimizer with β1 = 0.9, β2 = 0.98 and ε = 10^−9"
+    optimizer = torch.optim.Adam(
+        model.parameters(),
+        lr=1.0,  # LR is computed by the schedule, not a fixed value
+        betas=tuple(config["training"]["betas"]),
+        eps=config["training"]["eps"],
+    )
+
+    # §5.3, Eq. 3 — Custom LR schedule
+    scheduler = get_transformer_lr_schedule(
+        optimizer,
+        d_model=config["model"]["d_model"],
+        warmup_steps=config["training"]["warmup_steps"],
+    )
+
+    # §5.4 — Label smoothing with ε_ls = 0.1
+    loss_fn = LabelSmoothedCrossEntropy(
+        smoothing=config["training"]["label_smoothing"],
+    )
+
+    # --- Minimal forward/backward example with random data ---
+    model.train()
+    batch_size, src_len, tgt_len = 4, 20, 15
+
+    src = torch.randint(3, model_config.vocab_size, (batch_size, src_len), device=device)
+    tgt = torch.randint(3, model_config.vocab_size, (batch_size, tgt_len), device=device)
+    tgt_labels = torch.randint(3, model_config.vocab_size, (batch_size, tgt_len), device=device)
+
+    # Forward pass
+    optimizer.zero_grad()
+    logits = model(src, tgt)                             # (batch, tgt_len, vocab_size)
+    loss = loss_fn(logits, tgt_labels)
+    loss.backward()
+
+    # Gradient clipping — [UNSPECIFIED] paper does not mention gradient clipping
+    # Using: no clipping (following paper's silence on this)
+
+    optimizer.step()
+    scheduler.step()
+
+    current_lr = scheduler.get_last_lr()[0]
+    print(f"Step 1 | Loss: {loss.item():.4f} | LR: {current_lr:.2e}")
+    print(f"Logits shape: {logits.shape}")
+    print(f"\nTraining setup verified. For actual training, implement data loading from §5.1.")
+
+
+if __name__ == "__main__":
+    train()

plugins/codex/plugins/paper2code/skills/paper2code/worked/attention_is_all_you_need/src/utils.py

@@ -0,0 +1,96 @@
+"""
+Attention Is All You Need — Shared Utilities
+
+Paper: https://arxiv.org/abs/1706.03762
+Implements: Learning rate schedule (§5.3, Eq. 3) and masking utilities.
+"""
+
+import math
+
+import torch
+import torch.optim as optim
+
+
+def get_transformer_lr_schedule(
+    optimizer: optim.Optimizer,
+    d_model: int,
+    warmup_steps: int = 4000,
+) -> optim.lr_scheduler.LambdaLR:
+    """§5.3, Eq. 3 — Transformer learning rate schedule.
+
+    "We used the Adam optimizer with β1 = 0.9, β2 = 0.98 and ε = 10^−9.
+     We varied the learning rate over the course of training, according to the formula:
+
+         lr = d_model^(-0.5) * min(step_num^(-0.5), step_num * warmup_steps^(-1.5))
+
+     This corresponds to increasing the learning rate linearly for the first
+     warmup_steps training steps, and decreasing it thereafter proportionally
+     to the inverse square root of the step number."
+
+    Args:
+        optimizer: the optimizer to schedule
+        d_model: model dimension — §3, Table 3: 512
+        warmup_steps: §5.3 — "warmup_steps = 4000"
+
+    Returns:
+        LambdaLR scheduler (step per training step, not per epoch)
+    """
+    def lr_lambda(step: int) -> float:
+        step = max(step, 1)  # avoid division by zero at step 0
+        return d_model ** (-0.5) * min(step ** (-0.5), step * warmup_steps ** (-1.5))
+
+    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+def create_padding_mask(tokens: torch.Tensor, pad_idx: int = 0) -> torch.Tensor:
+    """Create a padding mask for source sequences.
+
+    Args:
+        tokens: (batch, seq_len) — token IDs
+        pad_idx: the padding token index
+
+    Returns:
+        (batch, 1, 1, seq_len) — True where tokens are NOT padding, False where padding
+    """
+    return (tokens != pad_idx).unsqueeze(1).unsqueeze(2)  # (batch, 1, 1, seq_len)
+
+
+def create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
+    """§3.2.3 — Create a causal (autoregressive) mask.
+
+    "We also modify the self-attention sub-layer in the decoder stack to prevent
+     positions from attending to subsequent positions."
+
+    Args:
+        seq_len: length of the target sequence
+        device: torch device
+
+    Returns:
+        (1, 1, seq_len, seq_len) — True for allowed positions, False for masked
+    """
+    mask = torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool))
+    return mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq_len, seq_len)
+
+
+def create_decoder_mask(
+    tgt: torch.Tensor, pad_idx: int = 0
+) -> torch.Tensor:
+    """Create combined causal + padding mask for decoder self-attention.
+
+    Args:
+        tgt: (batch, tgt_len) — target token IDs
+        pad_idx: padding token index
+
+    Returns:
+        (batch, 1, tgt_len, tgt_len) — combined mask
+    """
+    batch_size, tgt_len = tgt.size()
+
+    # Padding mask: (batch, 1, 1, tgt_len)
+    pad_mask = create_padding_mask(tgt, pad_idx)
+
+    # Causal mask: (1, 1, tgt_len, tgt_len)
+    causal_mask = create_causal_mask(tgt_len, tgt.device)
+
+    # Combine: both conditions must be True
+    return pad_mask & causal_mask  # (batch, 1, tgt_len, tgt_len)

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/README.md

@@ -0,0 +1,73 @@
+# Denoising Diffusion Probabilistic Models (DDPM)
+
+Implementation of [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239) (Ho, Jain, Abbeel, 2020).
+
+## What this implements
+
+The DDPM training and sampling procedure — a class of generative models that learn to reverse a gradual noising process. The core contribution is the training objective (simplified variational bound) and the reverse sampling algorithm, NOT the U-Net architecture (which is adapted from prior work). This implementation covers the forward diffusion process, the simplified training objective (L_simple), the noise schedule, and the reverse sampling algorithm from Algorithm 1 and Algorithm 2.
+
+## Quick start
+
+```bash
+pip install -r requirements.txt
+```
+
+```python
+from src.model import UNet, UNetConfig
+from src.loss import DDPMLoss
+from src.utils import linear_noise_schedule
+
+config = UNetConfig()
+model = UNet(config)
+noise_schedule = linear_noise_schedule(timesteps=1000)
+
+# Example: predict noise from noisy image at timestep t
+import torch
+x_t = torch.randn(2, 3, 32, 32)    # noisy image
+t = torch.randint(0, 1000, (2,))    # timesteps
+predicted_noise = model(x_t, t)
+print(predicted_noise.shape)  # (2, 3, 32, 32)
+```
+
+## File structure
+
+```
+ddpm/
+├── README.md                 # This file
+├── REPRODUCTION_NOTES.md     # Ambiguity audit — what's specified vs. assumed
+├── requirements.txt          # Dependencies
+├── src/
+│   ├── model.py              # U-Net noise prediction network (§3.3, Appendix B)
+│   ├── loss.py               # DDPM simplified loss L_simple (§3.4, Eq. 14)
+│   ├── data.py               # Dataset skeleton for image data
+│   ├── train.py              # Training loop — Algorithm 1
+│   ├── evaluate.py           # FID score computation
+│   └── utils.py              # Noise schedule, forward process, sampling (Algorithm 2)
+├── configs/
+│   └── base.yaml             # All hyperparameters from §4 and Appendix B
+└── notebooks/
+    └── walkthrough.ipynb     # Paper sections → code → sanity checks
+```
+
+## Important: Read REPRODUCTION_NOTES.md
+
+This implementation flags every choice that the paper does not specify.
+Before using this code for research, read [REPRODUCTION_NOTES.md](REPRODUCTION_NOTES.md)
+to understand which implementation details are from the paper and which are our choices.
+
+## Citation
+
+```bibtex
+@article{ho2020denoising,
+  title={Denoising diffusion probabilistic models},
+  author={Ho, Jonathan and Jain, Ajay and Abbeel, Pieter},
+  journal={Advances in neural information processing systems},
+  volume={33},
+  pages={6840--6851},
+  year={2020}
+}
+```
+
+---
+
+*Generated by [paper2code](https://github.com/PrathamLearnsToCode/paper2code) — citation-anchored paper implementation.*

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/REPRODUCTION_NOTES.md

@@ -0,0 +1,120 @@
+# Reproduction Notes: Denoising Diffusion Probabilistic Models
+
+> This document records every implementation choice, whether it was specified by the paper,
+> and what alternatives exist. If you're reproducing this paper, **read this first.**
+
+---
+
+## Paper
+
+- **Title:** Denoising Diffusion Probabilistic Models
+- **Authors:** Jonathan Ho, Ajay Jain, Pieter Abbeel
+- **Year:** 2020
+- **ArXiv:** https://arxiv.org/abs/2006.11239
+- **Official code:** https://github.com/hojonathanho/diffusion
+
+---
+
+## What this implements
+
+The DDPM training and sampling procedure: a generative model that learns to reverse a Markov chain of Gaussian noise additions. Given clean data x_0, the forward process gradually adds noise over T=1000 timesteps. The model (a U-Net) learns to predict the noise ε at each timestep, trained with the simplified objective L_simple = E[||ε - ε_θ(x_t, t)||²]. Sampling generates new images by iteratively denoising from pure Gaussian noise using Algorithm 2. The key contribution is the training objective and its connection to variational inference and score matching, NOT the U-Net architecture.
+
+---
+
+## Verified against
+
+- [x] Paper Algorithm 1 (Training) and Algorithm 2 (Sampling) — primary specification
+- [x] Paper §3.2, Eq. 11 — forward process q(x_t | x_0)
+- [x] Paper §3.4, Eq. 14 — simplified objective L_simple
+- [x] Paper §4 — noise schedule β_1 = 10^-4, β_T = 0.02
+- [x] Official code (github.com/hojonathanho/diffusion) — referenced for U-Net details
+- [x] Appendix B — model architecture details
+
+---
+
+## Unspecified choices
+
+| Component | Our Choice | Alternatives | Paper Quote (if partial) | Section |
+|-----------|-----------|--------------|--------------------------|---------|
+| U-Net channel multipliers | [1, 2, 2, 2] | [1, 2, 4, 8], varies | "channel multipliers" | App. B, Table after §B |
+| Attention resolution | 16 | 8, 16, 32 | "attention at 16×16 resolution" | App. B |
+| Dropout rate | 0.0 (CIFAR-10) | 0.1 (varies by dataset) | "dropout 0.0 for CIFAR10" | App. B |
+| Weight initialization | PyTorch defaults | Custom init from official code | — | — |
+| EMA start step | Step 0 | Various warmup strategies | — | — |
+| Gradient clipping | 1.0 | None | — | [FROM_OFFICIAL_CODE] github.com/hojonathanho/diffusion/blob/master/diffusion_tf/train_utils.py |
+| Learning rate warmup | None | Linear warmup for 5000 steps | — | — |
+| Number of residual blocks per level | 2 | 3, 4 | — | [FROM_OFFICIAL_CODE] |
+| Group norm num_groups | 32 | 8, 16 | — | [FROM_OFFICIAL_CODE] |
+| Sinusoidal timestep embedding dim | 128 | 256 | — | [FROM_OFFICIAL_CODE] |
+| Loss weighting | Uniform (L_simple) | SNR-based, learned | "We set the loss weight... to 1" | §3.4 |
+
+---
+
+## Known deviations
+
+| Deviation | Paper says | We do | Reason |
+|-----------|-----------|-------|--------|
+| Variance parameterization | §3.2 discusses learned Σ_θ but recommends fixed β_t For best results | We use fixed variance (σ²_t = β_t) | §3.4: "setting Σ_θ(x_t, t) = σ²_t I... achieved similar results" |
+| Variance choice | Paper discusses both σ²_t = β_t and σ²_t = β̃_t | We use σ²_t = β_t | §3.4: "We found that both had similar results" |
+
+---
+
+## Expected results
+
+| Metric | Paper's number | Dataset | Conditions |
+|--------|---------------|---------|------------|
+| FID | 3.17 | CIFAR-10 unconditional | Table 1, L_simple, 1000 timesteps |
+| IS | 9.46 | CIFAR-10 unconditional | Table 1, L_simple |
+| FID | 5.24 | LSUN Bedroom 256×256 | Table 1 |
+| NLL | ≤ 3.70 | CIFAR-10 | Table 2, bits/dim |
+
+**Note:** FID is sensitive to the number of generated samples (50k standard), the Inception
+model weights used, and random seeds. Reproducing exact FID numbers requires matching all of
+these. Using the `clean-fid` package is recommended for reproducible FID computation.
+
+---
+
+## Debugging tips
+
+1. **Samples look like noise after training**: Check that the noise schedule is correct. β should increase linearly from 1e-4 to 0.02. If α̅_T is not close to 0, the forward process doesn't destroy enough information and sampling from N(0,1) doesn't work.
+
+2. **Samples are blurry**: This is expected with small T or insufficient training. DDPM needs 800k+ training steps for good results on CIFAR-10. Also check that you're using EMA weights for sampling, not the training weights.
+
+3. **Mode collapse or repetitive samples**: Check that the timestep embedding is working correctly in the U-Net. If t is not properly conditioned, the model can't distinguish noise levels.
+
+4. **Loss doesn't decrease**: Verify the forward process q(x_t|x_0) is correct: x_t = √(α̅_t) * x_0 + √(1 - α̅_t) * ε. Check the sign and the square root.
+
+5. **NaN in sampling**: Check the variance computation at t=0. When t=0, some terms can be 0/0. The original code clips or special-cases t=0.
+
+---
+
+## Scope decisions
+
+### Implemented
+- Forward diffusion process q(x_t|x_0) (§3, Eq. 4) — core contribution
+- Training objective L_simple (§3.4, Eq. 14, Algorithm 1) — core contribution
+- Reverse sampling process (§3, Algorithm 2) — core contribution
+- Linear noise schedule β_1...β_T (§4) — required for the diffusion process
+- U-Net noise prediction network ε_θ (§3.3, Appendix B) — required model backbone
+- EMA of model weights (§4) — specified for evaluation
+
+### Intentionally excluded
+- Learned variance Σ_θ — paper finds fixed variance works similarly (§3.4)
+- Continuous-time diffusion — paper uses discrete timesteps
+- DDIM sampling — different paper (Song et al., 2020)
+- Classifier guidance — different paper (Dhariwal & Nichol, 2021)
+
+### Needed for full reproduction (not included)
+- CIFAR-10 dataset (torchvision can download it)
+- FID computation infrastructure (50k generated samples, Inception network)
+- 800k training steps on a single V100 GPU (~20 hours per DDPM paper)
+
+---
+
+## References
+
+- Ho et al., 2020 — primary paper
+- Sohl-Dickstein et al., 2015 — original diffusion framework
+- Ronneberger et al., 2015 (U-Net) — backbone architecture adapted in this work
+- Song & Ermon, 2019 — score matching perspective
+- Nichol & Dhariwal, 2021 (Improved DDPM) — improved noise schedule and learned variance

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/configs/base.yaml

@@ -0,0 +1,69 @@
+# Configuration for Denoising Diffusion Probabilistic Models (Ho et al., 2020)
+# Paper: https://arxiv.org/abs/2006.11239
+#
+# CONVENTION:
+#   - Lines with §X.Y cite the paper section specifying this value
+#   - Lines with [UNSPECIFIED] mark choices not stated in the paper
+#   - Lines with [FROM_OFFICIAL_CODE] are resolved from the authors' implementation
+#
+# All hyperparameters in one place. No magic numbers in code.
+
+# ---------------------------------------------------------------------------
+# Diffusion process — §2, §4
+# ---------------------------------------------------------------------------
+diffusion:
+  timesteps: 1000              # §4 — "T = 1000"
+  beta_start: 0.0001           # §4 — "β1 = 10^−4"
+  beta_end: 0.02               # §4 — "βT = 0.02"
+  schedule: "linear"           # §4 — "linearly from β1 = 10^-4 to βT = 0.02"
+  variance_type: "fixed_small" # §3.4 — "setting Σ_θ(x_t, t) = σ²_t I... achieved similar results"
+                               # "fixed_small" = β_t, "fixed_large" = β̃_t
+
+# ---------------------------------------------------------------------------
+# Model architecture — Appendix B, §3.3
+# ---------------------------------------------------------------------------
+model:
+  image_channels: 3            # §4 — RGB images
+  base_channels: 128           # Appendix B — "128 base channels" (CIFAR-10)
+  channel_mults: [1, 2, 2, 2]  # [FROM_OFFICIAL_CODE] — channel multipliers per resolution level
+  num_res_blocks: 2            # [FROM_OFFICIAL_CODE] — residual blocks per level
+  attention_resolutions: [16]  # Appendix B — "attention at 16×16 resolution"
+  dropout: 0.0                 # Appendix B — "dropout 0.0" (CIFAR-10). "dropout 0.1" for bedroom/church
+  time_embed_dim: 128          # [FROM_OFFICIAL_CODE] — sinusoidal timestep embedding dimension
+  num_groups: 32               # [FROM_OFFICIAL_CODE] — groups for GroupNorm
+
+# ---------------------------------------------------------------------------
+# Training — §4, Algorithm 1
+# ---------------------------------------------------------------------------
+training:
+  optimizer: "adam"             # §4 — "Adam"
+  lr: 0.0002                   # §4 — "learning rate 2 × 10^−4"
+  betas: [0.9, 0.999]          # [UNSPECIFIED] — Adam betas not stated, using PyTorch defaults
+  eps: 1.0e-8                  # [UNSPECIFIED] — Adam epsilon not stated
+  weight_decay: 0.0            # [UNSPECIFIED] — no weight decay mentioned
+  total_steps: 800000          # §4 — "We set the learning rate to 2 × 10^−4 without decay"
+                               # "800K iterations for CIFAR10" inferred from training time
+  batch_size: 128              # Appendix B — "batch size 128" (CIFAR-10)
+  gradient_clip: 1.0           # [FROM_OFFICIAL_CODE] — gradient clipping at 1.0
+  schedule: "constant"         # §4 — "without decay" (no LR schedule)
+
+  # EMA — §4
+  ema_decay: 0.9999            # §4 — "exponential moving average... decay factor of 0.9999"
+  ema_start_step: 0            # [UNSPECIFIED] — when to start EMA
+
+# ---------------------------------------------------------------------------
+# Data — §4
+# ---------------------------------------------------------------------------
+data:
+  dataset: "CIFAR-10"          # §4 — "CIFAR10"
+  image_size: 32               # CIFAR-10 images are 32×32
+  augmentation: "random_flip"  # §4 — "random horizontal flips"
+  normalize: true              # [UNSPECIFIED] — normalize images to [-1, 1] (standard for diffusion)
+
+# ---------------------------------------------------------------------------
+# Evaluation — §4
+# ---------------------------------------------------------------------------
+eval:
+  metric: "FID"                # §4 — "FID score"
+  num_samples: 50000           # Standard for FID computation
+  sampling_steps: 1000         # Same as T — full reverse process

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/notebooks/walkthrough.ipynb

@@ -0,0 +1,435 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "3569d851",
+   "metadata": {},
+   "source": [
+    "# Denoising Diffusion Probabilistic Models — Code Walkthrough\n",
+    "\n",
+    "**Paper**: Ho, Jain, Abbeel (2020) — [arXiv:2006.11239](https://arxiv.org/abs/2006.11239)\n",
+    "\n",
+    "This notebook walks through the core ideas and code of DDPM, connecting\n",
+    "each code fragment back to the paper's equations and algorithms.\n",
+    "\n",
+    "**All cells run on CPU** with tiny tensors for demonstration."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a4cb0c3b",
+   "metadata": {},
+   "source": [
+    "## 1. Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4431de7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.insert(0, '../src')\n",
+    "\n",
+    "import torch\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "from model import UNet, UNetConfig\n",
+    "from loss import DDPMLoss\n",
+    "from utils import linear_noise_schedule, q_sample, p_sample, sample\n",
+    "\n",
+    "torch.manual_seed(42)\n",
+    "device = torch.device('cpu')\n",
+    "print('Ready.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4ed2f605",
+   "metadata": {},
+   "source": [
+    "## 2. §2 — The Forward Process (Noise Schedule)\n",
+    "\n",
+    "From the paper:\n",
+    "\n",
+    "> \"We define the forward process... which produces latents $x_1, ..., x_T$\n",
+    "> by adding Gaussian noise according to a variance schedule $\\beta_1, ..., \\beta_T$\"\n",
+    "> — §2, Eq. 2\n",
+    "\n",
+    "$$q(x_t | x_{t-1}) = \\mathcal{N}(x_t; \\sqrt{1 - \\beta_t}\\, x_{t-1},\\; \\beta_t \\mathbf{I})$$\n",
+    "\n",
+    "The paper uses a **linear schedule** from $\\beta_1 = 10^{-4}$ to $\\beta_T = 0.02$ with $T = 1000$ (§4)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b245ab62",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "T = 1000\n",
+    "betas = linear_noise_schedule(T, beta_start=1e-4, beta_end=0.02)\n",
+    "alphas = 1.0 - betas\n",
+    "alpha_bar = torch.cumprod(alphas, dim=0)\n",
+    "\n",
+    "print(f'β_1 = {betas[0]:.6f}, β_T = {betas[-1]:.6f}')\n",
+    "print(f'ᾱ_1 = {alpha_bar[0]:.6f}, ᾱ_T = {alpha_bar[-1]:.6f}')\n",
+    "\n",
+    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
+    "axes[0].plot(betas.numpy())\n",
+    "axes[0].set_title('Linear β schedule (§4)')\n",
+    "axes[0].set_xlabel('Timestep t')\n",
+    "axes[0].set_ylabel('β_t')\n",
+    "\n",
+    "axes[1].plot(alpha_bar.numpy())\n",
+    "axes[1].set_title('Cumulative ᾱ_t = ∏α_s (Eq. 4)')\n",
+    "axes[1].set_xlabel('Timestep t')\n",
+    "axes[1].set_ylabel('ᾱ_t')\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c0c5e659",
+   "metadata": {},
+   "source": [
+    "## 3. §2, Eq. 4 — Closed-Form Forward Process\n",
+    "\n",
+    "A key property: we can sample $x_t$ directly from $x_0$ without iterating:\n",
+    "\n",
+    "$$q(x_t | x_0) = \\mathcal{N}(x_t; \\sqrt{\\bar{\\alpha}_t}\\, x_0,\\; (1 - \\bar{\\alpha}_t) \\mathbf{I})$$\n",
+    "\n",
+    "This is implemented in `q_sample()`. Let's visualize the noising process."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "894eb3d8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create a tiny fake image (3×32×32)\n",
+    "x_0 = torch.randn(1, 3, 32, 32)  # pretend this is a CIFAR-10 image\n",
+    "\n",
+    "sqrt_alpha_bar = torch.sqrt(alpha_bar)\n",
+    "sqrt_one_minus_alpha_bar = torch.sqrt(1.0 - alpha_bar)\n",
+    "\n",
+    "# Visualize noising at different timesteps\n",
+    "timesteps = [0, 100, 300, 500, 700, 999]\n",
+    "fig, axes = plt.subplots(1, len(timesteps), figsize=(18, 3))\n",
+    "\n",
+    "for i, t_val in enumerate(timesteps):\n",
+    "    t = torch.tensor([t_val])\n",
+    "    noise = torch.randn_like(x_0)\n",
+    "    x_t = q_sample(x_0, t, sqrt_alpha_bar, sqrt_one_minus_alpha_bar, noise)\n",
+    "    \n",
+    "    # Show first channel\n",
+    "    img = x_t[0, 0].numpy()\n",
+    "    axes[i].imshow(img, cmap='gray', vmin=-3, vmax=3)\n",
+    "    axes[i].set_title(f't={t_val}\\nᾱ={alpha_bar[t_val]:.3f}')\n",
+    "    axes[i].axis('off')\n",
+    "\n",
+    "plt.suptitle('Forward process: q(x_t | x_0) — Eq. 4', fontsize=14)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "00dc6e6e",
+   "metadata": {},
+   "source": [
+    "## 4. §3.3 — U-Net Architecture\n",
+    "\n",
+    "From the paper:\n",
+    "\n",
+    "> \"The neural network... uses a U-Net backbone similar to an unmasked\n",
+    "> PixelCNN++ with group normalization throughout, and... self-attention\n",
+    "> at the 16×16 feature map resolution.\" — §3.3\n",
+    "\n",
+    "Let's instantiate a **tiny** version for CPU testing."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "92b1c172",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Tiny config for CPU testing (real config uses base_channels=128)\n",
+    "tiny_config = UNetConfig(\n",
+    "    image_channels=3,\n",
+    "    base_channels=16,          # Real: 128 (Appendix B)\n",
+    "    channel_mults=(1, 2),      # Real: (1, 2, 2, 2)\n",
+    "    num_res_blocks=1,          # Real: 2\n",
+    "    attention_resolutions=(16,),\n",
+    "    dropout=0.0,\n",
+    "    num_groups=8,              # Real: 32\n",
+    "    image_size=32,\n",
+    ")\n",
+    "\n",
+    "tiny_model = UNet(tiny_config)\n",
+    "print(tiny_model)\n",
+    "\n",
+    "# Sanity check: forward pass\n",
+    "x = torch.randn(2, 3, 32, 32)  # (batch=2, C=3, H=32, W=32)\n",
+    "t = torch.randint(1, T + 1, (2,))  # (batch=2,)\n",
+    "noise_pred = tiny_model(x, t)\n",
+    "print(f'\\nInput shape:  {x.shape}')\n",
+    "print(f'Output shape: {noise_pred.shape}')\n",
+    "assert x.shape == noise_pred.shape, 'Output must match input shape!'\n",
+    "print('✓ Shape check passed')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3cdbe712",
+   "metadata": {},
+   "source": [
+    "## 5. §3.4, Eq. 14 — The Simplified Loss (L_simple)\n",
+    "\n",
+    "From the paper:\n",
+    "\n",
+    "> \"$L_{\\text{simple}} = \\mathbb{E}_{t, \\mathbf{x}_0, \\boldsymbol{\\epsilon}}\\left[\\|\\boldsymbol{\\epsilon} - \\boldsymbol{\\epsilon}_\\theta(\\sqrt{\\bar{\\alpha}_t}\\mathbf{x}_0 + \\sqrt{1 - \\bar{\\alpha}_t}\\boldsymbol{\\epsilon}, t)\\|^2\\right]$\"\n",
+    "> — §3.4, Eq. 14\n",
+    "\n",
+    "This is simply MSE between the true noise $\\epsilon$ and the predicted noise $\\epsilon_\\theta$."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15c1e63a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "criterion = DDPMLoss()\n",
+    "\n",
+    "# Simulate one training step (Algorithm 1)\n",
+    "x_0 = torch.randn(4, 3, 32, 32)  # batch of 4 fake images\n",
+    "t = torch.randint(1, T + 1, (4,))\n",
+    "noise = torch.randn_like(x_0)\n",
+    "\n",
+    "# Forward process: get x_t\n",
+    "x_t = q_sample(x_0, t, sqrt_alpha_bar, sqrt_one_minus_alpha_bar, noise)\n",
+    "\n",
+    "# Model predicts noise\n",
+    "noise_pred = tiny_model(x_t, t)\n",
+    "\n",
+    "# Loss: MSE(ε, ε_θ)\n",
+    "loss = criterion(noise_pred, noise)\n",
+    "print(f'L_simple = {loss.item():.4f}')\n",
+    "print(f'Expected: ~1.0 for untrained model (predicting random noise)')\n",
+    "\n",
+    "# Verify gradient flows\n",
+    "loss.backward()\n",
+    "grad_norm = sum(p.grad.norm().item() for p in tiny_model.parameters() if p.grad is not None)\n",
+    "print(f'Total gradient norm: {grad_norm:.4f}')\n",
+    "print('✓ Gradients flow correctly')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "89e3cdea",
+   "metadata": {},
+   "source": [
+    "## 6. Algorithm 1 — One Training Step\n",
+    "\n",
+    "Let's verify the complete training loop on tiny data:\n",
+    "\n",
+    "```\n",
+    "Algorithm 1 Training\n",
+    "1: repeat\n",
+    "2:   x_0 ~ q(x_0)\n",
+    "3:   t ~ Uniform({1,...,T})\n",
+    "4:   ε ~ N(0, I)\n",
+    "5:   Take gradient step on ∇_θ ||ε − ε_θ(√ᾱ_t x_0 + √(1−ᾱ_t) ε, t)||²\n",
+    "6: until converged\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "22e7ffef",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "optimizer = torch.optim.Adam(tiny_model.parameters(), lr=2e-4)  # §B — lr=2e-4\n",
+    "losses = []\n",
+    "\n",
+    "tiny_model.train()\n",
+    "# Overfit on a single batch to verify training works\n",
+    "fixed_x0 = torch.randn(8, 3, 32, 32)\n",
+    "\n",
+    "for step in range(50):\n",
+    "    # Line 3: t ~ Uniform({1,...,T})\n",
+    "    t = torch.randint(1, T + 1, (8,))\n",
+    "    # Line 4: ε ~ N(0, I) \n",
+    "    noise = torch.randn_like(fixed_x0)\n",
+    "    # Compute x_t\n",
+    "    x_t = q_sample(fixed_x0, t, sqrt_alpha_bar, sqrt_one_minus_alpha_bar, noise)\n",
+    "    # Predict noise\n",
+    "    noise_pred = tiny_model(x_t, t)\n",
+    "    # Line 5: gradient step\n",
+    "    loss = criterion(noise_pred, noise)\n",
+    "    optimizer.zero_grad()\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "    losses.append(loss.item())\n",
+    "\n",
+    "plt.plot(losses)\n",
+    "plt.xlabel('Step')\n",
+    "plt.ylabel('L_simple')\n",
+    "plt.title('Algorithm 1 — Training loss (tiny model, overfitting)')\n",
+    "plt.show()\n",
+    "\n",
+    "print(f'Initial loss: {losses[0]:.4f}')\n",
+    "print(f'Final loss:   {losses[-1]:.4f}')\n",
+    "print('✓ Loss decreases — training works')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "51f8754d",
+   "metadata": {},
+   "source": [
+    "## 7. Algorithm 2 — Sampling (Reverse Process)\n",
+    "\n",
+    "From the paper:\n",
+    "\n",
+    "```\n",
+    "Algorithm 2 Sampling\n",
+    "1: x_T ~ N(0, I)\n",
+    "2: for t = T, ..., 1 do\n",
+    "3:   z ~ N(0, I) if t > 1, else z = 0\n",
+    "4:   x_{t-1} = (1/√α_t)(x_t − (β_t/√(1−ᾱ_t)) ε_θ(x_t, t)) + σ_t z\n",
+    "5: end for\n",
+    "6: return x_0\n",
+    "```\n",
+    "\n",
+    "The key is Eq. 11: the mean of $p_\\theta(x_{t-1}|x_t)$ is parameterized via the noise prediction."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "58f95564",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Use tiny T for CPU demo\n",
+    "T_demo = 20\n",
+    "betas_demo = linear_noise_schedule(T_demo, 1e-4, 0.02)\n",
+    "\n",
+    "tiny_model.eval()\n",
+    "\n",
+    "# Sample 4 images\n",
+    "with torch.no_grad():\n",
+    "    samples = sample(tiny_model, shape=(4, 3, 32, 32), T=T_demo,\n",
+    "                     betas=betas_demo, device=device)\n",
+    "\n",
+    "print(f'Sample shape: {samples.shape}')\n",
+    "print(f'Sample range: [{samples.min():.2f}, {samples.max():.2f}]')\n",
+    "\n",
+    "fig, axes = plt.subplots(1, 4, figsize=(12, 3))\n",
+    "for i in range(4):\n",
+    "    img = (samples[i].permute(1, 2, 0).numpy() + 1) / 2  # [-1,1] -> [0,1]\n",
+    "    img = img.clip(0, 1)\n",
+    "    axes[i].imshow(img)\n",
+    "    axes[i].axis('off')\n",
+    "    axes[i].set_title(f'Sample {i+1}')\n",
+    "plt.suptitle('Algorithm 2 — Sampled images (untrained tiny model)', fontsize=14)\n",
+    "plt.tight_layout()\n",
+    "plt.show()\n",
+    "print('✓ Sampling pipeline works (images are noise — model is untrained)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "97397882",
+   "metadata": {},
+   "source": [
+    "## 8. Noise Prediction Verification\n",
+    "\n",
+    "Let's verify that our `p_sample` correctly implements Eq. 11:\n",
+    "\n",
+    "$$\\mu_\\theta(x_t, t) = \\frac{1}{\\sqrt{\\alpha_t}}\\left(x_t - \\frac{\\beta_t}{\\sqrt{1 - \\bar{\\alpha}_t}} \\epsilon_\\theta(x_t, t)\\right)$$"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fa855792",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Verify the math of p_sample matches Eq. 11\n",
+    "alphas_demo = 1.0 - betas_demo\n",
+    "alpha_bar_demo = torch.cumprod(alphas_demo, dim=0)\n",
+    "\n",
+    "x_t = torch.randn(1, 3, 32, 32)\n",
+    "t_val = 10\n",
+    "t_tensor = torch.tensor([t_val])\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    eps_pred = tiny_model(x_t, t_tensor)\n",
+    "\n",
+    "# Manual computation of μ_θ (Eq. 11)\n",
+    "alpha_t = alphas_demo[t_val]            # α_t\n",
+    "alpha_bar_t = alpha_bar_demo[t_val]     # ᾱ_t  \n",
+    "beta_t = betas_demo[t_val]              # β_t\n",
+    "\n",
+    "mu_manual = (1.0 / torch.sqrt(alpha_t)) * (\n",
+    "    x_t - (beta_t / torch.sqrt(1.0 - alpha_bar_t)) * eps_pred\n",
+    ")\n",
+    "\n",
+    "# Verify mean computation matches\n",
+    "print(f'α_t = {alpha_t:.6f}')\n",
+    "print(f'ᾱ_t = {alpha_bar_t:.6f}')\n",
+    "print(f'β_t = {beta_t:.6f}')\n",
+    "print(f'μ_θ mean value: {mu_manual.mean():.6f}')\n",
+    "print(f'μ_θ shape: {mu_manual.shape}')\n",
+    "print('✓ Eq. 11 computation verified')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dd36cbb7",
+   "metadata": {},
+   "source": [
+    "## 9. Summary\n",
+    "\n",
+    "This notebook verified the core components of the DDPM implementation:\n",
+    "\n",
+    "| Component | Section | Status |\n",
+    "|---|---|---|\n",
+    "| Linear noise schedule | §4 | ✓ |\n",
+    "| Forward process q(x_t\\|x_0) | §2, Eq. 4 | ✓ |\n",
+    "| U-Net ε_θ(x_t, t) | §3.3 | ✓ |\n",
+    "| L_simple loss | §3.4, Eq. 14 | ✓ |\n",
+    "| Algorithm 1 (Training) | §3.4 | ✓ |\n",
+    "| Algorithm 2 (Sampling) | §3.4 | ✓ |\n",
+    "| Eq. 11 mean computation | §3.2 | ✓ |\n",
+    "\n",
+    "**To train for real**: Use `train.py` with the full config (128 base channels,\n",
+    "800K steps, GPU) and evaluate with `evaluate.py`.\n",
+    "\n",
+    "**Expected FID**: ~3.17 with full training (Table 1)."
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.0.0
+numpy>=1.24.0
+pyyaml>=6.0

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/review.md

@@ -0,0 +1,100 @@
+# Review — DDPM (2006.11239) Implementation
+
+## Summary
+This is a minimal, citation-anchored implementation of Denoising Diffusion
+Probabilistic Models (Ho, Jain, Abbeel 2020) for unconditional image generation
+on CIFAR-10.
+
+---
+
+## Coverage Assessment
+
+### What the paper contributes (§1, §3)
+| Contribution | Implemented? | Notes |
+|---|---|---|
+| Forward diffusion process q(x_t\|x_{t-1}) | ✓ | `utils.py: q_sample()` — §2, Eq. 2 |
+| Reverse process p_θ(x_{t-1}\|x_t) | ✓ | `utils.py: p_sample()` — §3.2, Eq. 11 |
+| L_simple training objective | ✓ | `loss.py: DDPMLoss` — §3.4, Eq. 14 |
+| Linear noise schedule | ✓ | `utils.py: linear_noise_schedule()` — §4 |
+| Algorithm 1 (Training) | ✓ | `train.py` — §3.4 |
+| Algorithm 2 (Sampling) | ✓ | `utils.py: sample()` — §3.4 |
+| U-Net architecture | ✓ | `model.py: UNet` — §3.3, Appendix B |
+| Fixed variance σ²_t = β_t | ✓ | §3.4 recommendation |
+| EMA of model parameters | ✓ | `utils.py: EMA` — §4 |
+
+### What is NOT implemented
+| Paper content | Reason for exclusion |
+|---|---|
+| L_vlb (full variational bound, Eq. 5) | Paper recommends L_simple; L_vlb used only for comparison in Table 2 |
+| Learned variance σ²_t | Paper explores this but recommends fixed variance (§3.4) |
+| Log-likelihood evaluation (Table 2) | Not the primary contribution; FID is the main metric |
+| CelebA-HQ / LSUN experiments | Larger-scale experiments; architecture is CIFAR-10 config |
+| Progressive coding / lossless compression (§4.3) | Interesting analysis but not the generative model itself |
+| Interpolation experiments (§4.2) | Evaluation analysis, not implementation |
+
+---
+
+## Faithfulness to Paper
+
+### ✓ Correctly implemented
+- **Noise schedule**: Linear β_1 = 10^-4 to β_T = 0.02 for T = 1000 (§4)
+- **Training objective**: Pure MSE between ε and ε_θ, matching Eq. 14
+- **Forward process**: Closed-form q(x_t|x_0) via reparameterization (Eq. 4)
+- **Reverse process**: Fixed variance, learned mean via Eq. 11
+- **U-Net**: GroupNorm, self-attention at 16×16, residual blocks (§3.3, Appendix B)
+- **EMA decay 0.9999** (§4)
+
+### ⚠ Assumptions and gaps
+- **Batch size 128**: [FROM_OFFICIAL_CODE] — paper does not explicitly state batch size
+- **Gradient clipping 1.0**: [FROM_OFFICIAL_CODE] — not mentioned in paper
+- **No learning rate schedule**: [UNSPECIFIED] — paper doesn't describe any schedule
+- **U-Net channel multipliers [1,2,2,2]**: [FROM_OFFICIAL_CODE] — Appendix B only says "128 channels"
+- **Data normalization [-1,1]**: [FROM_OFFICIAL_CODE] — paper doesn't specify range
+- **Zero-init for output conv**: [FROM_OFFICIAL_CODE]
+
+### ✗ Known differences from ideal
+- **Model size**: Our U-Net is approximately correct for the CIFAR-10 config but
+  may not match the exact parameter count (~35.7M reported in Appendix B)
+- **Attention implementation**: We use a simple single-head self-attention; the
+  official code has minor differences in normalization
+- **No multi-GPU training**: Paper trained on TPU v3-8; we provide single-GPU code
+- **800K steps**: Paper trains for 800K steps which takes ~50 GPU-hours on V100
+
+---
+
+## Reproducibility Expectations
+
+### CIFAR-10 (Unconditional, 32×32)
+
+| Metric | Paper (Table 1) | Expected from this code |
+|---|---|---|
+| FID | 3.17 | ~5-10 (due to potential minor arch differences) |
+| IS | 9.46 | ~8-9 |
+
+Achieving exact paper numbers requires:
+1. Exact architecture match (some details only in official code)
+2. Full 800K training steps
+3. EMA parameters for evaluation
+4. 50K generated samples for FID
+
+---
+
+## Code Quality Notes
+
+1. **Citation anchoring**: Every function and class references the paper section
+   where its specification originates
+2. **Flag discipline**: All assumptions/unofficial-code-sourced details are marked
+   with [FROM_OFFICIAL_CODE], [UNSPECIFIED], or [ASSUMPTION]
+3. **Shape comments**: Tensor operations include shape annotations
+4. **No unnecessary code**: No class-conditional generation, no classifier guidance,
+   no DDIM sampling — these are all later work, not in this paper
+
+---
+
+## What to improve if extending
+
+1. **DDIM sampling** (Song et al., 2020): Faster sampling with fewer steps
+2. **Classifier-free guidance** (Ho & Salimans, 2022): Better conditional generation
+3. **Cosine schedule** (Nichol & Dhariwal, 2021): Better sample quality
+4. **Learned variance** (Nichol & Dhariwal, 2021): Avoids the fixed variance assumption
+5. **Multi-GPU training**: Use PyTorch DDP for faster training

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/src/data.py

@@ -0,0 +1,111 @@
+"""
+Denoising Diffusion Probabilistic Models — Data Loading
+
+Paper: https://arxiv.org/abs/2006.11239
+Authors: Ho, Jain, Abbeel (2020)
+
+§4 — "We set T = 1000... on CIFAR10 (32×32)... We used random horizontal
+flips during training; we tried training both with and without dropout..."
+
+Data preprocessing:
+  - Normalize images to [-1, 1] range (not [0,1])
+    [FROM_OFFICIAL_CODE] — Paper does not state normalization range explicitly,
+    but the official code normalizes to [-1,1] which is consistent with
+    the Gaussian output distribution N(x_0; μ_θ, σ²I).
+  - Random horizontal flips (§4)
+  - No other augmentation mentioned
+"""
+
+from typing import Optional, Tuple
+
+import torch
+from torch.utils.data import DataLoader, Dataset
+from torchvision import datasets, transforms
+
+
+def get_cifar10_transforms(image_size: int = 32) -> Tuple[transforms.Compose, transforms.Compose]:
+    """Return train and test transforms for CIFAR-10.
+
+    §4 — "random horizontal flips during training"
+    [FROM_OFFICIAL_CODE] — normalize to [-1, 1]
+
+    Args:
+        image_size: Target image size (CIFAR-10 is natively 32×32)
+
+    Returns:
+        (train_transform, test_transform)
+    """
+    train_transform = transforms.Compose([
+        transforms.RandomHorizontalFlip(),       # §4 — explicit mention
+        transforms.ToTensor(),                    # [0, 255] -> [0, 1]
+        transforms.Normalize(                     # [0, 1] -> [-1, 1]
+            mean=[0.5, 0.5, 0.5],
+            std=[0.5, 0.5, 0.5],
+        ),
+    ])
+
+    test_transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.5, 0.5, 0.5],
+            std=[0.5, 0.5, 0.5],
+        ),
+    ])
+
+    return train_transform, test_transform
+
+
+def get_dataloaders(
+    data_dir: str = "./data",
+    batch_size: int = 128,
+    num_workers: int = 4,
+    image_size: int = 32,
+) -> Tuple[DataLoader, DataLoader]:
+    """Create CIFAR-10 train and test dataloaders.
+
+    §4 — CIFAR-10 unconditional generation at 32×32
+    §4 — "batch size 128" [FROM_OFFICIAL_CODE — paper does not state batch size]
+
+    Args:
+        data_dir: Root directory for dataset download/cache
+        batch_size: Batch size for training
+        num_workers: Number of data loading workers
+        image_size: Image resolution (32 for CIFAR-10)
+
+    Returns:
+        (train_loader, test_loader)
+    """
+    train_transform, test_transform = get_cifar10_transforms(image_size)
+
+    train_dataset = datasets.CIFAR10(
+        root=data_dir,
+        train=True,
+        download=True,
+        transform=train_transform,
+    )
+
+    test_dataset = datasets.CIFAR10(
+        root=data_dir,
+        train=False,
+        download=True,
+        transform=test_transform,
+    )
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=True,
+        drop_last=True,  # [ASSUMPTION] drop incomplete last batch for stable training
+    )
+
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=True,
+    )
+
+    return train_loader, test_loader

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/src/evaluate.py

@@ -0,0 +1,246 @@
+"""
+Denoising Diffusion Probabilistic Models — Evaluation (FID Score)
+
+Paper: https://arxiv.org/abs/2006.11239
+Authors: Ho, Jain, Abbeel (2020)
+
+§4 — "We report FID score and Inception Score... Our best results
+are FID: 3.17, IS: 9.46 (unconditional) on CIFAR-10."
+
+This module provides a wrapper for sampling from a trained DDPM model
+and computing FID scores using the pytorch-fid library.
+
+FID (Fréchet Inception Distance) is the primary evaluation metric
+used in §4 and Table 1. Lower FID = better quality.
+
+NOTE: FID computation requires:
+  1. Generating 50K samples (standard for CIFAR-10)
+  2. Computing Inception features for real and generated images
+  3. Computing the Fréchet distance between the two feature distributions
+
+This is computationally expensive. For quick validation, generate a
+small batch and visually inspect.
+"""
+
+import os
+import logging
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torchvision
+import yaml
+
+from model import UNet, UNetConfig
+from utils import linear_noise_schedule, sample, EMA
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s — %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def load_model(
+    checkpoint_path: str,
+    device: torch.device,
+    use_ema: bool = True,
+) -> tuple:
+    """Load a trained DDPM model from checkpoint.
+
+    §4 — "we also report results with an exponential moving average"
+    The EMA parameters typically produce better samples.
+
+    Args:
+        checkpoint_path: Path to .pt checkpoint file
+        device: Target device
+        use_ema: Whether to load EMA parameters (recommended)
+
+    Returns:
+        (model, config_dict)
+    """
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    cfg = checkpoint["config"]
+    model_cfg = cfg["model"]
+
+    unet_config = UNetConfig(
+        image_channels=model_cfg.get("image_channels", 3),
+        base_channels=model_cfg.get("base_channels", 128),
+        channel_mults=tuple(model_cfg.get("channel_mults", [1, 2, 2, 2])),
+        num_res_blocks=model_cfg.get("num_res_blocks", 2),
+        attention_resolutions=tuple(model_cfg.get("attention_resolutions", [16])),
+        dropout=model_cfg.get("dropout", 0.0),
+        num_groups=model_cfg.get("num_groups", 32),
+        image_size=cfg["data"].get("image_size", 32),
+    )
+
+    model = UNet(unet_config).to(device)
+
+    if use_ema and "ema_state_dict" in checkpoint:
+        # Load EMA parameters
+        ema_params = checkpoint["ema_state_dict"]
+        for name, param in model.named_parameters():
+            if name in ema_params:
+                param.data.copy_(ema_params[name])
+        logger.info("Loaded EMA parameters")
+    else:
+        model.load_state_dict(checkpoint["model_state_dict"])
+        logger.info("Loaded model parameters (no EMA)")
+
+    return model, cfg
+
+
+@torch.no_grad()
+def generate_samples(
+    model: UNet,
+    config: dict,
+    num_samples: int = 64,
+    device: Optional[torch.device] = None,
+) -> torch.Tensor:
+    """Generate samples using Algorithm 2 — Sampling.
+
+    §3.4 — "sampling from p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), σ²_t I)"
+
+    Args:
+        model: Trained UNet model
+        config: Config dict from checkpoint
+        num_samples: Number of images to generate
+        device: Target device
+
+    Returns:
+        (num_samples, C, H, W) — generated images in [0, 1] range
+    """
+    if device is None:
+        device = next(model.parameters()).device
+
+    diff_cfg = config["diffusion"]
+    data_cfg = config["data"]
+
+    T = diff_cfg["T"]
+    betas = linear_noise_schedule(T, diff_cfg["beta_start"], diff_cfg["beta_end"]).to(device)
+
+    image_size = data_cfg.get("image_size", 32)
+    image_channels = config["model"].get("image_channels", 3)
+    shape = (num_samples, image_channels, image_size, image_size)
+
+    model.eval()
+    samples = sample(model, shape, T, betas, device)
+
+    # Convert from [-1, 1] to [0, 1] for saving
+    samples = (samples + 1.0) / 2.0
+    samples = samples.clamp(0.0, 1.0)
+
+    return samples
+
+
+def save_samples(
+    samples: torch.Tensor,
+    output_dir: str,
+    prefix: str = "sample",
+    make_grid: bool = True,
+    nrow: int = 8,
+):
+    """Save generated samples as images.
+
+    Args:
+        samples: (N, C, H, W) in [0, 1]
+        output_dir: Directory to save images
+        prefix: Filename prefix
+        make_grid: If True, also save a grid image
+        nrow: Number of images per row in grid
+    """
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    if make_grid:
+        grid = torchvision.utils.make_grid(samples, nrow=nrow, padding=2)
+        grid_path = output_dir / f"{prefix}_grid.png"
+        torchvision.utils.save_image(grid, grid_path)
+        logger.info(f"Saved grid: {grid_path}")
+
+    # Save individual images
+    for i, img in enumerate(samples):
+        img_path = output_dir / f"{prefix}_{i:05d}.png"
+        torchvision.utils.save_image(img, img_path)
+
+    logger.info(f"Saved {len(samples)} individual images to {output_dir}")
+
+
+def compute_fid(
+    generated_dir: str,
+    real_stats_path: Optional[str] = None,
+    batch_size: int = 50,
+    device: str = "cuda",
+    dims: int = 2048,
+) -> float:
+    """Compute FID score between generated samples and real data.
+
+    §4 — "We report FID score... Our best results are FID: 3.17"
+
+    Requires the pytorch-fid package: pip install pytorch-fid
+
+    For CIFAR-10, you need pre-computed stats for the real training set,
+    or provide a directory of real images.
+
+    Args:
+        generated_dir: Directory containing generated .png images
+        real_stats_path: Path to pre-computed .npz stats for real data,
+                         OR directory containing real images
+        batch_size: Batch size for Inception feature extraction
+        device: Device for computation
+        dims: Inception feature dimensionality (2048 = pool3)
+
+    Returns:
+        FID score (float). Lower is better.
+    """
+    try:
+        from pytorch_fid import fid_score
+    except ImportError:
+        logger.error(
+            "pytorch-fid not installed. Install with: pip install pytorch-fid\n"
+            "Then re-run evaluation."
+        )
+        raise
+
+    if real_stats_path is None:
+        raise ValueError(
+            "Must provide real_stats_path: either a .npz file with pre-computed "
+            "Inception statistics, or a directory of real CIFAR-10 images."
+        )
+
+    fid = fid_score.calculate_fid_given_paths(
+        [generated_dir, real_stats_path],
+        batch_size=batch_size,
+        device=torch.device(device),
+        dims=dims,
+    )
+
+    logger.info(f"FID score: {fid:.2f}")
+    return fid
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="DDPM Evaluation — Generate samples and compute FID")
+    parser.add_argument("--checkpoint", type=str, required=True,
+                        help="Path to trained model checkpoint")
+    parser.add_argument("--num_samples", type=int, default=64,
+                        help="Number of samples to generate")
+    parser.add_argument("--output_dir", type=str, default="generated",
+                        help="Output directory for generated images")
+    parser.add_argument("--fid", action="store_true",
+                        help="Compute FID score (requires --real_stats)")
+    parser.add_argument("--real_stats", type=str, default=None,
+                        help="Path to real data stats (.npz) or directory")
+    parser.add_argument("--no_ema", action="store_true",
+                        help="Don't use EMA parameters")
+    parser.add_argument("--device", type=str, default="cuda",
+                        help="Device (cuda/cpu)")
+    args = parser.parse_args()
+
+    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+
+    model, cfg = load_model(args.checkpoint, device, use_ema=not args.no_ema)
+    samples = generate_samples(model, cfg, args.num_samples, device)
+    save_samples(samples, args.output_dir)
+
+    if args.fid:
+        compute_fid(args.output_dir, args.real_stats, device=args.device)

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/src/loss.py

@@ -0,0 +1,60 @@
+"""
+Denoising Diffusion Probabilistic Models — Training Objective (Simplified Loss)
+
+Paper: https://arxiv.org/abs/2006.11239
+Authors: Ho, Jain, Abbeel (2020)
+
+Implements: L_simple from §3.4, Eq. 14
+
+"We have shown that the variational bound... can be optimized with a
+ simplified objective... which resembles denoising score matching..."
+
+  L_simple = E_{t, x_0, ε} [ ||ε − ε_θ(√ᾱ_t x_0 + √(1−ᾱ_t) ε, t)||² ]
+
+This is MSE between the true noise ε and the model's prediction ε_θ.
+The expectation is over:
+  - t ~ Uniform({1, ..., T})
+  - x_0 ~ q(x_0)  (data distribution)
+  - ε ~ N(0, I)
+
+§3.4 — "We found it beneficial to sample quality (and simpler to implement)
+to train on the following variant of the variational bound... L_simple"
+
+§3.4 — "Algorithm 1" describes the training procedure that uses this loss.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class DDPMLoss(nn.Module):
+    """§3.4, Eq. 14 — Simplified training objective L_simple.
+
+    Computes MSE between true noise and predicted noise:
+        L = ||ε − ε_θ(x_t, t)||²
+
+    This module handles only the loss computation. The caller (training loop)
+    is responsible for sampling t, computing x_t from x_0, and calling the model.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(
+        self,
+        noise_pred: torch.Tensor,
+        noise_true: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        §3.4, Eq. 14 — L_simple = E[||ε − ε_θ(x_t, t)||²]
+
+        Args:
+            noise_pred: (batch, C, H, W) — predicted noise ε_θ(x_t, t)
+            noise_true: (batch, C, H, W) — true noise ε ~ N(0, I)
+
+        Returns:
+            scalar — mean squared error loss
+        """
+        # §3.4 — Simple MSE between predicted and true noise
+        # "equivalent to (a re-weighted variant of) the ELBO"
+        return nn.functional.mse_loss(noise_pred, noise_true)

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/src/model.py

@@ -0,0 +1,382 @@
+"""
+Denoising Diffusion Probabilistic Models — U-Net Noise Prediction Network
+
+Paper: https://arxiv.org/abs/2006.11239
+Authors: Ho, Jain, Abbeel (2020)
+
+Implements: U-Net architecture for noise prediction ε_θ(x_t, t)
+The U-Net is adapted from PixelCNN++ (Salimans et al., 2017) and the architecture
+used in score matching (Song & Ermon, 2019). It is NOT the paper's core contribution
+but is required as the backbone for the diffusion process.
+
+Section references:
+  §3.3 — "The neural network... is similar to an unmasked PixelCNN++ with
+          group normalization and self-attention"
+  Appendix B — Architecture details (channel counts, attention resolutions)
+
+NOTE: This U-Net implementation follows the architecture from the official code
+(github.com/hojonathanho/diffusion) since the paper describes it only briefly
+in §3.3 and Appendix B. Many details are [FROM_OFFICIAL_CODE].
+"""
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+@dataclass
+class UNetConfig:
+    """U-Net configuration.
+
+    Values from Appendix B (CIFAR-10 config) and official code unless marked.
+    """
+    image_channels: int = 3           # §4 — RGB images
+    base_channels: int = 128          # Appendix B — "128 base channels"
+    channel_mults: tuple = (1, 2, 2, 2)  # [FROM_OFFICIAL_CODE] — channel multipliers per level
+    num_res_blocks: int = 2           # [FROM_OFFICIAL_CODE] — residual blocks per resolution level
+    attention_resolutions: tuple = (16,)  # Appendix B — "attention at 16×16 resolution"
+    dropout: float = 0.0             # Appendix B — "dropout 0.0" for CIFAR-10
+    time_embed_dim: int = 512        # [FROM_OFFICIAL_CODE] — 4 * base_channels
+    num_groups: int = 32             # [FROM_OFFICIAL_CODE] — groups for GroupNorm
+    image_size: int = 32             # CIFAR-10 is 32×32
+
+
+# ---------------------------------------------------------------------------
+# Time embedding — sinusoidal (borrowed from Transformer positional encoding)
+# ---------------------------------------------------------------------------
+
+class SinusoidalTimeEmbedding(nn.Module):
+    """Sinusoidal timestep embedding, following the Transformer positional encoding.
+
+    §3.3 — "The diffusion time t is specified by adding the Transformer
+    sinusoidal position embedding into each residual block."
+
+    [FROM_OFFICIAL_CODE] The embedding dimension and MLP projection follow
+    the official implementation.
+    """
+
+    def __init__(self, embed_dim: int):
+        super().__init__()
+        self.embed_dim = embed_dim
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            t: (batch,) — integer timesteps
+
+        Returns:
+            (batch, embed_dim) — sinusoidal embedding
+        """
+        half_dim = self.embed_dim // 2
+        emb = math.log(10000.0) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=t.device, dtype=torch.float32) * -emb)
+        emb = t.float().unsqueeze(1) * emb.unsqueeze(0)  # (batch, half_dim)
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)  # (batch, embed_dim)
+        return emb
+
+
+# ---------------------------------------------------------------------------
+# Building blocks
+# ---------------------------------------------------------------------------
+
+class ResidualBlock(nn.Module):
+    """Residual block with time embedding conditioning.
+
+    §3.3 — "group normalization throughout... Transformer sinusoidal position
+    embedding into each residual block"
+
+    [FROM_OFFICIAL_CODE] Structure: GroupNorm -> SiLU -> Conv -> GroupNorm -> SiLU -> Dropout -> Conv + residual
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        time_embed_dim: int,
+        dropout: float = 0.0,
+        num_groups: int = 32,
+    ):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(num_groups, in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+
+        # Time embedding projection
+        self.time_proj = nn.Linear(time_embed_dim, out_channels)
+
+        self.norm2 = nn.GroupNorm(num_groups, out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+
+        # Skip connection (1x1 conv if channel count changes)
+        if in_channels != out_channels:
+            self.skip = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+        else:
+            self.skip = nn.Identity()
+
+    def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (batch, in_channels, H, W)
+            t_emb: (batch, time_embed_dim)
+
+        Returns:
+            (batch, out_channels, H, W)
+        """
+        h = self.norm1(x)
+        h = F.silu(h)                                  # (batch, in_channels, H, W)
+        h = self.conv1(h)                              # (batch, out_channels, H, W)
+
+        # Add time embedding
+        t = self.time_proj(F.silu(t_emb))              # (batch, out_channels)
+        h = h + t.unsqueeze(-1).unsqueeze(-1)          # (batch, out_channels, H, W) — broadcast
+
+        h = self.norm2(h)
+        h = F.silu(h)
+        h = self.dropout(h)
+        h = self.conv2(h)                              # (batch, out_channels, H, W)
+
+        return h + self.skip(x)                         # residual connection
+
+
+class AttentionBlock(nn.Module):
+    """Self-attention block for the U-Net.
+
+    §3.3 — "self-attention... at the 16×16 feature map resolution"
+    Appendix B — "We add one head of self-attention at the 16×16 resolution"
+
+    [FROM_OFFICIAL_CODE] Uses a single attention head with GroupNorm.
+    """
+
+    def __init__(self, channels: int, num_groups: int = 32):
+        super().__init__()
+        self.norm = nn.GroupNorm(num_groups, channels)
+        self.qkv = nn.Conv1d(channels, channels * 3, kernel_size=1)
+        self.proj = nn.Conv1d(channels, channels, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (batch, channels, H, W)
+        Returns:
+            (batch, channels, H, W)
+        """
+        batch, channels, h, w = x.shape
+        residual = x
+
+        x = self.norm(x)
+        x = x.view(batch, channels, h * w)        # (batch, channels, H*W)
+
+        qkv = self.qkv(x)                          # (batch, 3*channels, H*W)
+        q, k, v = qkv.chunk(3, dim=1)              # each: (batch, channels, H*W)
+
+        # Scaled dot-product attention
+        scale = 1.0 / math.sqrt(channels)
+        attn = torch.bmm(q.transpose(1, 2), k) * scale  # (batch, H*W, H*W)
+        attn = F.softmax(attn, dim=-1)
+
+        out = torch.bmm(v, attn.transpose(1, 2))   # (batch, channels, H*W)
+        out = self.proj(out)                         # (batch, channels, H*W)
+        out = out.view(batch, channels, h, w)        # (batch, channels, H, W)
+
+        return out + residual
+
+
+class Downsample(nn.Module):
+    """Spatial downsampling by factor 2. [FROM_OFFICIAL_CODE]"""
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)  # (batch, C, H, W) -> (batch, C, H/2, W/2)
+
+
+class Upsample(nn.Module):
+    """Spatial upsampling by factor 2. [FROM_OFFICIAL_CODE]"""
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=2, mode="nearest")
+        return self.conv(x)  # (batch, C, H, W) -> (batch, C, 2H, 2W)
+
+
+# ---------------------------------------------------------------------------
+# §3.3 — Full U-Net
+# ---------------------------------------------------------------------------
+
+class UNet(nn.Module):
+    """§3.3, Appendix B — U-Net noise prediction network ε_θ(x_t, t).
+
+    "We use a U-Net backbone similar to an unmasked PixelCNN++ with
+     group normalization throughout, and we add one head of self-attention
+     at the 16×16 feature map resolution."
+
+    The U-Net takes a noisy image x_t and a timestep t, and predicts the
+    noise ε that was added. This is NOT the paper's core contribution —
+    it is the backbone model that enables the diffusion process.
+
+    Architecture (for CIFAR-10 32×32):
+      Down: 32→32→16→8→4 (with skip connections)
+      Middle: bottleneck with attention
+      Up: 4→8→16→32→32 (with skip connections from down path)
+    """
+
+    def __init__(self, config: UNetConfig):
+        super().__init__()
+        self.config = config
+        ch = config.base_channels
+
+        # Time embedding: sinusoidal -> MLP
+        # §3.3 — "Transformer sinusoidal position embedding"
+        time_embed_dim = config.time_embed_dim
+        self.time_embed = nn.Sequential(
+            SinusoidalTimeEmbedding(ch),
+            nn.Linear(ch, time_embed_dim),
+            nn.SiLU(),
+            nn.Linear(time_embed_dim, time_embed_dim),
+        )
+
+        # Initial convolution
+        self.input_conv = nn.Conv2d(config.image_channels, ch, kernel_size=3, padding=1)
+
+        # Downsampling path
+        self.down_blocks = nn.ModuleList()
+        self.down_samples = nn.ModuleList()
+        channels = [ch]
+        current_res = config.image_size
+        in_ch = ch
+
+        for level, mult in enumerate(config.channel_mults):
+            out_ch = ch * mult
+            for _ in range(config.num_res_blocks):
+                layers = [ResidualBlock(in_ch, out_ch, time_embed_dim,
+                                        config.dropout, config.num_groups)]
+                if current_res in config.attention_resolutions:
+                    layers.append(AttentionBlock(out_ch, config.num_groups))
+                self.down_blocks.append(nn.ModuleList(layers))
+                channels.append(out_ch)
+                in_ch = out_ch
+
+            if level < len(config.channel_mults) - 1:
+                self.down_samples.append(Downsample(out_ch))
+                channels.append(out_ch)
+                current_res //= 2
+            else:
+                self.down_samples.append(nn.Identity())
+
+        # Middle (bottleneck)
+        self.mid_block1 = ResidualBlock(in_ch, in_ch, time_embed_dim,
+                                         config.dropout, config.num_groups)
+        self.mid_attn = AttentionBlock(in_ch, config.num_groups)
+        self.mid_block2 = ResidualBlock(in_ch, in_ch, time_embed_dim,
+                                         config.dropout, config.num_groups)
+
+        # Upsampling path
+        self.up_blocks = nn.ModuleList()
+        self.up_samples = nn.ModuleList()
+
+        for level in reversed(range(len(config.channel_mults))):
+            mult = config.channel_mults[level]
+            out_ch = ch * mult
+            for i in range(config.num_res_blocks + 1):
+                skip_ch = channels.pop()
+                layers = [ResidualBlock(in_ch + skip_ch, out_ch, time_embed_dim,
+                                        config.dropout, config.num_groups)]
+                if current_res in config.attention_resolutions:
+                    layers.append(AttentionBlock(out_ch, config.num_groups))
+                self.up_blocks.append(nn.ModuleList(layers))
+                in_ch = out_ch
+
+            if level > 0:
+                self.up_samples.append(Upsample(out_ch))
+                current_res *= 2
+            else:
+                self.up_samples.append(nn.Identity())
+
+        # Output
+        self.output_norm = nn.GroupNorm(config.num_groups, in_ch)
+        self.output_conv = nn.Conv2d(in_ch, config.image_channels, kernel_size=3, padding=1)
+
+        # [UNSPECIFIED] Zero-initialize the final conv (from official code)
+        nn.init.zeros_(self.output_conv.weight)
+        nn.init.zeros_(self.output_conv.bias)
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """Predict noise ε_θ(x_t, t).
+
+        Args:
+            x: (batch, C, H, W) — noisy image x_t
+            t: (batch,) — integer timesteps
+
+        Returns:
+            (batch, C, H, W) — predicted noise ε_θ
+        """
+        # Time embedding
+        t_emb = self.time_embed(t)  # (batch, time_embed_dim)
+
+        # Initial conv
+        h = self.input_conv(x)  # (batch, base_channels, H, W)
+
+        # Downsampling with skip connections
+        skips = [h]
+        block_idx = 0
+        for level in range(len(self.config.channel_mults)):
+            for _ in range(self.config.num_res_blocks):
+                layers = self.down_blocks[block_idx]
+                h = layers[0](h, t_emb)  # ResidualBlock
+                if len(layers) > 1:
+                    h = layers[1](h)      # AttentionBlock (if present)
+                skips.append(h)
+                block_idx += 1
+
+            h = self.down_samples[level](h)
+            if not isinstance(self.down_samples[level], nn.Identity):
+                skips.append(h)
+
+        # Middle
+        h = self.mid_block1(h, t_emb)
+        h = self.mid_attn(h)
+        h = self.mid_block2(h, t_emb)
+
+        # Upsampling with skip connections
+        block_idx = 0
+        for level in reversed(range(len(self.config.channel_mults))):
+            for _ in range(self.config.num_res_blocks + 1):
+                skip = skips.pop()
+                h = torch.cat([h, skip], dim=1)  # Concatenate skip connection
+                layers = self.up_blocks[block_idx]
+                h = layers[0](h, t_emb)  # ResidualBlock
+                if len(layers) > 1:
+                    h = layers[1](h)      # AttentionBlock (if present)
+                block_idx += 1
+
+            h = self.up_samples[level - len(self.config.channel_mults)](h) if level > 0 else h
+
+        # Output
+        h = self.output_norm(h)
+        h = F.silu(h)
+        return self.output_conv(h)  # (batch, C, H, W) — predicted noise
+
+    def __repr__(self) -> str:
+        total_params = sum(p.numel() for p in self.parameters())
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        return (
+            f"UNet(\n"
+            f"  image_size={self.config.image_size}, base_channels={self.config.base_channels},\n"
+            f"  channel_mults={self.config.channel_mults}, num_res_blocks={self.config.num_res_blocks},\n"
+            f"  total_params={total_params:,},\n"
+            f"  trainable_params={trainable_params:,}\n"
+            f")"
+        )

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/src/train.py

@@ -0,0 +1,195 @@
+"""
+Denoising Diffusion Probabilistic Models — Training Loop
+
+Paper: https://arxiv.org/abs/2006.11239
+Authors: Ho, Jain, Abbeel (2020)
+
+Implements: Algorithm 1 — Training procedure
+
+  1: repeat
+  2:   x_0 ~ q(x_0)                         ← sample data
+  3:   t ~ Uniform({1, ..., T})              ← sample timestep
+  4:   ε ~ N(0, I)                           ← sample noise
+  5:   Take gradient step on
+          ∇_θ ||ε − ε_θ(√ᾱ_t x_0 + √(1−ᾱ_t) ε, t)||²
+  6: until converged
+
+Hyperparameters from §4 and Appendix B:
+  - Adam optimizer, learning rate 2e-4 (§B)
+  - No lr schedule/warmup [UNSPECIFIED — not mentioned in paper]
+  - Gradient clipping at 1.0 [FROM_OFFICIAL_CODE]
+  - EMA with decay 0.9999 (§4)
+  - T = 1000 (§4)
+  - Trained for 800K steps on CIFAR-10 (Appendix B)
+"""
+
+import os
+import logging
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import yaml
+
+from model import UNet, UNetConfig
+from loss import DDPMLoss
+from data import get_dataloaders
+from utils import linear_noise_schedule, q_sample, EMA
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s — %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def train(config_path: str = "configs/base.yaml"):
+    """Algorithm 1 — DDPM Training.
+
+    Args:
+        config_path: Path to YAML config file
+    """
+    # --- Load config ---
+    config_path = Path(config_path)
+    if config_path.exists():
+        with open(config_path) as f:
+            cfg = yaml.safe_load(f)
+    else:
+        raise FileNotFoundError(f"Config not found: {config_path}")
+
+    diff_cfg = cfg["diffusion"]
+    model_cfg = cfg["model"]
+    train_cfg = cfg["training"]
+    data_cfg = cfg["data"]
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"Using device: {device}")
+
+    # --- Noise schedule ---
+    # §2, Eq. 4 — linear schedule β_1 = 0.0001, β_T = 0.02
+    T = diff_cfg["T"]
+    betas = linear_noise_schedule(T, diff_cfg["beta_start"], diff_cfg["beta_end"])
+    betas = betas.to(device)
+
+    alphas = 1.0 - betas                                    # α_t = 1 − β_t
+    alpha_bar = torch.cumprod(alphas, dim=0)                # ᾱ_t = ∏_{s=1}^{t} α_s
+    sqrt_alpha_bar = torch.sqrt(alpha_bar)                  # √ᾱ_t
+    sqrt_one_minus_alpha_bar = torch.sqrt(1.0 - alpha_bar)  # √(1−ᾱ_t)
+
+    # --- Model ---
+    unet_config = UNetConfig(
+        image_channels=model_cfg.get("image_channels", 3),
+        base_channels=model_cfg.get("base_channels", 128),
+        channel_mults=tuple(model_cfg.get("channel_mults", [1, 2, 2, 2])),
+        num_res_blocks=model_cfg.get("num_res_blocks", 2),
+        attention_resolutions=tuple(model_cfg.get("attention_resolutions", [16])),
+        dropout=model_cfg.get("dropout", 0.0),
+        num_groups=model_cfg.get("num_groups", 32),
+        image_size=data_cfg.get("image_size", 32),
+    )
+    model = UNet(unet_config).to(device)
+    logger.info(f"Model: {model}")
+
+    # --- EMA ---
+    # §4 — "we also report results with an exponential moving average of
+    # model parameters with a decay factor of 0.9999"
+    ema = EMA(model, decay=train_cfg.get("ema_decay", 0.9999))
+
+    # --- Optimizer ---
+    # Appendix B — "Adam, lr = 2 × 10^-4"
+    optimizer = torch.optim.Adam(
+        model.parameters(),
+        lr=float(train_cfg.get("lr", 2e-4)),
+    )
+
+    # --- Loss ---
+    criterion = DDPMLoss()
+
+    # --- Data ---
+    train_loader, _ = get_dataloaders(
+        data_dir=data_cfg.get("data_dir", "./data"),
+        batch_size=train_cfg.get("batch_size", 128),
+        num_workers=data_cfg.get("num_workers", 4),
+        image_size=data_cfg.get("image_size", 32),
+    )
+
+    # --- Training loop: Algorithm 1 ---
+    total_steps = train_cfg.get("total_steps", 800_000)
+    log_every = train_cfg.get("log_every", 1000)
+    save_every = train_cfg.get("save_every", 50_000)
+    save_dir = Path(train_cfg.get("save_dir", "checkpoints"))
+    save_dir.mkdir(parents=True, exist_ok=True)
+    grad_clip = train_cfg.get("gradient_clip", 1.0)
+
+    step = 0
+    model.train()
+
+    while step < total_steps:
+        for batch in train_loader:
+            if step >= total_steps:
+                break
+
+            # Algorithm 1, line 2: x_0 ~ q(x_0)
+            x_0 = batch[0].to(device)  # (batch, C, H, W), labels discarded (unconditional)
+            batch_size = x_0.shape[0]
+
+            # Algorithm 1, line 3: t ~ Uniform({1, ..., T})
+            t = torch.randint(1, T + 1, (batch_size,), device=device)
+
+            # Algorithm 1, line 4: ε ~ N(0, I)
+            noise = torch.randn_like(x_0)
+
+            # Algorithm 1, line 5: compute x_t and predict noise
+            x_t = q_sample(x_0, t, sqrt_alpha_bar, sqrt_one_minus_alpha_bar, noise)
+            noise_pred = model(x_t, t)
+
+            # L_simple — §3.4, Eq. 14
+            loss = criterion(noise_pred, noise)
+
+            # Algorithm 1, line 5: gradient step
+            optimizer.zero_grad()
+            loss.backward()
+
+            # [FROM_OFFICIAL_CODE] gradient clipping
+            if grad_clip > 0:
+                nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+
+            optimizer.step()
+
+            # §4 — EMA update
+            ema.update()
+
+            step += 1
+
+            if step % log_every == 0:
+                logger.info(f"Step {step}/{total_steps} — loss: {loss.item():.6f}")
+
+            if step % save_every == 0:
+                checkpoint = {
+                    "step": step,
+                    "model_state_dict": model.state_dict(),
+                    "ema_state_dict": ema.shadow_params,
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "loss": loss.item(),
+                    "config": cfg,
+                }
+                ckpt_path = save_dir / f"ddpm_step_{step}.pt"
+                torch.save(checkpoint, ckpt_path)
+                logger.info(f"Saved checkpoint: {ckpt_path}")
+
+    # Final save
+    final_path = save_dir / "ddpm_final.pt"
+    torch.save({
+        "step": step,
+        "model_state_dict": model.state_dict(),
+        "ema_state_dict": ema.shadow_params,
+        "optimizer_state_dict": optimizer.state_dict(),
+        "config": cfg,
+    }, final_path)
+    logger.info(f"Training complete. Final checkpoint: {final_path}")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="Train DDPM — Algorithm 1")
+    parser.add_argument("--config", type=str, default="configs/base.yaml",
+                        help="Path to config YAML")
+    args = parser.parse_args()
+    train(args.config)

plugins/codex/plugins/paper2code/skills/paper2code/worked/ddpm/src/utils.py

@@ -0,0 +1,243 @@
+"""
+Denoising Diffusion Probabilistic Models — Shared Utilities
+
+Paper: https://arxiv.org/abs/2006.11239
+Implements: Noise schedule (§4), forward process q(x_t|x_0) (§3, Eq. 4),
+            and reverse sampling (Algorithm 2).
+
+Section references:
+  §2 — Forward and reverse diffusion processes
+  §3 — Training objective derivation
+  §4 — Noise schedule specification
+  Algorithm 1 — Training
+  Algorithm 2 — Sampling
+"""
+
+import math
+from typing import Dict, Optional
+
+import torch
+import torch.nn as nn
+
+
+def linear_noise_schedule(
+    timesteps: int = 1000,
+    beta_start: float = 1e-4,
+    beta_end: float = 0.02,
+) -> Dict[str, torch.Tensor]:
+    """§4 — Linear variance schedule.
+
+    "We set the forward process variances to constants increasing linearly
+     from β_1 = 10^-4 to β_T = 0.02."
+
+    Precomputes all quantities needed for training and sampling.
+
+    Args:
+        timesteps: §4 — T = 1000
+        beta_start: §4 — β_1 = 10^-4
+        beta_end: §4 — β_T = 0.02
+
+    Returns:
+        Dict with precomputed schedule tensors
+    """
+    # §4 — linear schedule
+    betas = torch.linspace(beta_start, beta_end, timesteps)
+
+    # §2 — α_t = 1 - β_t
+    alphas = 1.0 - betas
+
+    # §2 — α̅_t = Π_{s=1}^{t} α_s  (cumulative product)
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+
+    # α̅_{t-1} — needed for posterior q(x_{t-1} | x_t, x_0)
+    alphas_cumprod_prev = torch.cat([torch.tensor([1.0]), alphas_cumprod[:-1]])
+
+    # Precompute quantities for q(x_t | x_0) = N(x_t; √α̅_t x_0, (1-α̅_t)I)
+    sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
+    sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)
+
+    # Precompute for sampling (Algorithm 2)
+    sqrt_recip_alphas = 1.0 / torch.sqrt(alphas)
+
+    # §3.2, Eq. 7 — Posterior variance β̃_t = β_t * (1 - α̅_{t-1}) / (1 - α̅_t)
+    posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+
+    return {
+        "betas": betas,
+        "alphas": alphas,
+        "alphas_cumprod": alphas_cumprod,
+        "alphas_cumprod_prev": alphas_cumprod_prev,
+        "sqrt_alphas_cumprod": sqrt_alphas_cumprod,
+        "sqrt_one_minus_alphas_cumprod": sqrt_one_minus_alphas_cumprod,
+        "sqrt_recip_alphas": sqrt_recip_alphas,
+        "posterior_variance": posterior_variance,
+    }
+
+
+def q_sample(
+    x_0: torch.Tensor,
+    t: torch.Tensor,
+    schedule: Dict[str, torch.Tensor],
+    noise: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """§3, Eq. 4 — Forward process: sample x_t from q(x_t | x_0).
+
+    "A notable property is that we can sample x_t at any arbitrary time step t
+     in closed form: q(x_t | x_0) = N(x_t; √α̅_t x_0, (1 - α̅_t)I)"
+
+    x_t = √α̅_t * x_0 + √(1 - α̅_t) * ε, where ε ~ N(0, I)
+
+    Args:
+        x_0: (batch, C, H, W) — clean images
+        t: (batch,) — timestep indices
+        schedule: precomputed noise schedule
+        noise: optional pre-sampled noise (for reproducibility)
+
+    Returns:
+        x_t: (batch, C, H, W) — noisy images at timestep t
+    """
+    if noise is None:
+        noise = torch.randn_like(x_0)
+
+    # Extract schedule values for timestep t, reshape for broadcasting
+    sqrt_alpha_cumprod = schedule["sqrt_alphas_cumprod"][t]          # (batch,)
+    sqrt_one_minus_alpha_cumprod = schedule["sqrt_one_minus_alphas_cumprod"][t]  # (batch,)
+
+    # Reshape for broadcasting with (batch, C, H, W)
+    sqrt_alpha_cumprod = sqrt_alpha_cumprod.view(-1, 1, 1, 1)
+    sqrt_one_minus_alpha_cumprod = sqrt_one_minus_alpha_cumprod.view(-1, 1, 1, 1)
+
+    # §3, Eq. 4 — x_t = √α̅_t * x_0 + √(1 - α̅_t) * ε
+    return sqrt_alpha_cumprod * x_0 + sqrt_one_minus_alpha_cumprod * noise
+
+
+@torch.no_grad()
+def p_sample(
+    model: nn.Module,
+    x_t: torch.Tensor,
+    t: torch.Tensor,
+    t_index: int,
+    schedule: Dict[str, torch.Tensor],
+) -> torch.Tensor:
+    """Algorithm 2, lines 3-4 — Single reverse step: sample x_{t-1} from p_θ(x_{t-1} | x_t).
+
+    "x_{t-1} = 1/√α_t * (x_t - β_t/√(1-α̅_t) * ε_θ(x_t, t)) + σ_t * z"
+
+    where z ~ N(0, I) if t > 1, else z = 0.
+
+    Args:
+        model: noise prediction network ε_θ
+        x_t: (batch, C, H, W) — current noisy sample
+        t: (batch,) — current timestep (as tensor for model input)
+        t_index: integer timestep (for indexing schedule)
+        schedule: precomputed noise schedule
+
+    Returns:
+        x_{t-1}: (batch, C, H, W) — denoised sample one step
+    """
+    # Predict noise ε_θ(x_t, t)
+    predicted_noise = model(x_t, t)  # (batch, C, H, W)
+
+    # Extract schedule values
+    beta_t = schedule["betas"][t_index]
+    sqrt_recip_alpha_t = schedule["sqrt_recip_alphas"][t_index]
+    sqrt_one_minus_alpha_cumprod_t = schedule["sqrt_one_minus_alphas_cumprod"][t_index]
+
+    # Algorithm 2, line 4 — Compute mean of p_θ(x_{t-1} | x_t)
+    # μ_θ = 1/√α_t * (x_t - β_t/√(1-α̅_t) * ε_θ(x_t, t))
+    mean = sqrt_recip_alpha_t * (
+        x_t - beta_t / sqrt_one_minus_alpha_cumprod_t * predicted_noise
+    )
+
+    if t_index == 0:
+        # Algorithm 2, line 3 — z = 0 when t = 1 (final step)
+        return mean
+    else:
+        # Algorithm 2, line 3 — z ~ N(0, I) when t > 1
+        # §3.4 — σ²_t = β_t (fixed small variance)
+        sigma_t = torch.sqrt(schedule["betas"][t_index])
+        noise = torch.randn_like(x_t)
+        return mean + sigma_t * noise
+
+
+@torch.no_grad()
+def sample(
+    model: nn.Module,
+    schedule: Dict[str, torch.Tensor],
+    image_shape: tuple,
+    device: torch.device,
+) -> torch.Tensor:
+    """Algorithm 2 — Full reverse sampling process.
+
+    "Algorithm 2 Sampling
+     1: x_T ~ N(0, I)
+     2: for t = T, ..., 1 do
+     3:   z ~ N(0, I) if t > 1, else z = 0
+     4:   x_{t-1} = 1/√α_t (x_t - β_t/√(1-α̅_t) ε_θ(x_t, t)) + σ_t z
+     5: end for
+     6: return x_0"
+
+    Args:
+        model: noise prediction network ε_θ (should be in eval mode, ideally EMA weights)
+        schedule: precomputed noise schedule
+        image_shape: (batch, C, H, W) — shape of images to generate
+        device: torch device
+
+    Returns:
+        x_0: (batch, C, H, W) — generated images
+    """
+    model.eval()
+    timesteps = len(schedule["betas"])
+
+    # Algorithm 2, line 1: x_T ~ N(0, I)
+    x = torch.randn(image_shape, device=device)
+
+    # Algorithm 2, lines 2-5: reverse iterate from t=T to t=1
+    for t_index in reversed(range(timesteps)):
+        t = torch.full((image_shape[0],), t_index, device=device, dtype=torch.long)
+        x = p_sample(model, x, t, t_index, schedule)
+
+    return x
+
+
+class EMA:
+    """§4 — Exponential Moving Average of model parameters.
+
+    "We report sample quality metrics using an exponential moving average (EMA)
+     of model parameters with a decay factor of 0.9999."
+
+    The EMA weights are used for sampling/evaluation, not for training.
+    """
+
+    def __init__(self, model: nn.Module, decay: float = 0.9999):
+        """
+        Args:
+            model: the model to track
+            decay: §4 — "decay factor of 0.9999"
+        """
+        self.decay = decay
+        self.shadow = {}
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                self.shadow[name] = param.data.clone()
+
+    @torch.no_grad()
+    def update(self, model: nn.Module):
+        """Update EMA weights after each training step."""
+        for name, param in model.named_parameters():
+            if param.requires_grad and name in self.shadow:
+                self.shadow[name].mul_(self.decay).add_(
+                    param.data, alpha=1.0 - self.decay
+                )
+
+    def apply(self, model: nn.Module):
+        """Load EMA weights into model (for evaluation/sampling)."""
+        for name, param in model.named_parameters():
+            if name in self.shadow:
+                param.data.copy_(self.shadow[name])
+
+    def restore(self, model: nn.Module):
+        """Restore original model weights (after evaluation)."""
+        # NOTE: This requires storing original weights separately.
+        # The caller should save model.state_dict() before calling apply().
+        pass