Enhance Colab badge functionality: Update script to add 'Open in Colab' badges to both template and solution notebooks. Modify badge generation to accommodate different notebook directories and ensure proper processing of both templates and solutions.

scripts/add_colab_badges.py

@@ -1,27 +1,29 @@
 #!/usr/bin/env python3
-"""Add 'Open in Colab' badges to all template notebooks."""
+"""Add 'Open in Colab' badges to all template and solution notebooks."""
 
 import json
 from pathlib import Path
 
 REPO = "duoan/TorchCode"
 BRANCH = "master"
-TEMPLATES_DIR = Path(__file__).resolve().parent.parent / "templates"
+ROOT = Path(__file__).resolve().parent.parent
+TEMPLATES_DIR = ROOT / "templates"
+SOLUTIONS_DIR = ROOT / "solutions"
 BADGE_IMG = "https://colab.research.google.com/assets/colab-badge.svg"
 
 
-def colab_url(filename: str) -> str:
+def colab_url(filename: str, folder: str) -> str:
     return (
         f"https://colab.research.google.com/github/{REPO}"
-        f"/blob/{BRANCH}/templates/{filename}"
+        f"/blob/{BRANCH}/{folder}/{filename}"
     )
 
 
-def badge_markdown(filename: str) -> str:
-    return f"[![Open In Colab]({BADGE_IMG})]({colab_url(filename)})"
+def badge_markdown(filename: str, folder: str) -> str:
+    return f"[![Open In Colab]({BADGE_IMG})]({colab_url(filename, folder)})"
 
 
-def process_notebook(path: Path) -> bool:
+def process_notebook(path: Path, folder: str) -> bool:
     with open(path, "r", encoding="utf-8") as f:
         nb = json.load(f)
 
@@ -34,7 +36,7 @@ def process_notebook(path: Path) -> bool:
     if "colab-badge.svg" in flat:
         return False
 
-    badge = badge_markdown(path.name)
+    badge = badge_markdown(path.name, folder)
     cells[0]["source"] = [badge + "\n\n"] + (
         source_lines if isinstance(source_lines, list) else [source_lines]
     )
@@ -49,11 +51,17 @@ def process_notebook(path: Path) -> bool:
 def main() -> None:
     updated = 0
     for nb_path in sorted(TEMPLATES_DIR.glob("*.ipynb")):
-        if process_notebook(nb_path):
-            print(f"  ✅ {nb_path.name}")
+        if process_notebook(nb_path, "templates"):
+            print(f"  ✅ templates/{nb_path.name}")
             updated += 1
         else:
-            print(f"  ⏭️  {nb_path.name} (already has badge or skipped)")
+            print(f"  ⏭️  templates/{nb_path.name} (already has badge or skipped)")
+    for nb_path in sorted(SOLUTIONS_DIR.glob("*.ipynb")):
+        if process_notebook(nb_path, "solutions"):
+            print(f"  ✅ solutions/{nb_path.name}")
+            updated += 1
+        else:
+            print(f"  ⏭️  solutions/{nb_path.name} (already has badge or skipped)")
     print(f"\nDone — updated {updated} notebooks.")
 
 

scripts/add_colab_torch_judge_install.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Add Colab-only pip install of torch-judge to all notebooks that use torch_judge."""
+
+import json
+from pathlib import Path
+
+INSTALL_CELL_SOURCE = [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n",
+]
+
+MARKER = "get_ipython().run_line_magic('pip', 'install"
+
+
+def has_torch_judge(nb: dict) -> bool:
+    for cell in nb.get("cells", []):
+        src = cell.get("source", [])
+        flat = "".join(src) if isinstance(src, list) else str(src)
+        if "torch_judge" in flat:
+            return True
+    return False
+
+
+def already_has_install(nb: dict) -> bool:
+    for cell in nb.get("cells", []):
+        src = cell.get("source", [])
+        flat = "".join(src) if isinstance(src, list) else str(src)
+        if MARKER in flat and "torch-judge" in flat:
+            return True
+    return False
+
+
+def process_notebook(path: Path) -> bool:
+    with open(path, "r", encoding="utf-8") as f:
+        nb = json.load(f)
+
+    if not has_torch_judge(nb):
+        return False
+    if already_has_install(nb):
+        return False
+
+    cells = nb["cells"]
+    if not cells:
+        return False
+
+    # Insert install cell at index 1 (after first cell, usually markdown title)
+    install_cell = {
+        "cell_type": "code",
+        "metadata": {},
+        "source": INSTALL_CELL_SOURCE,
+        "outputs": [],
+        "execution_count": None,
+    }
+    cells.insert(1, install_cell)
+
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(nb, f, ensure_ascii=False, indent=1)
+        f.write("\n")
+
+    return True
+
+
+def main() -> None:
+    root = Path(__file__).resolve().parent.parent
+    updated = 0
+    for pattern in ["templates/*.ipynb", "solutions/*.ipynb"]:
+        for path in sorted(root.glob(pattern)):
+            if process_notebook(path):
+                print(f"  + {path.relative_to(root)}")
+                updated += 1
+    print(f"Updated {updated} notebooks.")
+
+
+if __name__ == "__main__":
+    main()

solutions/01_relu_solution.ipynb

@@ -1,73 +1,89 @@
 {
- "nbformat": 4,
- "nbformat_minor": 5,
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.11.0"
-  }
- },
  "cells": [
   {
    "cell_type": "markdown",
+   "id": "0556419b",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/01_relu_solution.ipynb)\n",
+    "\n",
     "# 🟢 Solution: Implement ReLU\n",
     "\n",
     "Reference solution for the ReLU activation function.\n",
     "\n",
     "$$\\text{ReLU}(x) = \\max(0, x)$$"
-   ],
-   "outputs": []
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
-    "import torch"
-   ],
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
    "outputs": [],
-   "execution_count": null
+   "source": [
+    "import torch"
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "# ✅ SOLUTION\n",
     "\n",
     "def relu(x: torch.Tensor) -> torch.Tensor:\n",
     "    return x * (x > 0).float()"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "# Verify\n",
     "x = torch.tensor([-2., -1., 0., 1., 2.])\n",
     "print(\"Input: \", x)\n",
     "print(\"Output:\", relu(x))"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "# Run judge\n",
     "from torch_judge import check\n",
     "check(\"relu\")"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   }
- ]
-}
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

solutions/02_softmax_solution.ipynb

@@ -17,6 +17,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/02_softmax_solution.ipynb)\n\n",
     "# 🟢 Solution: Implement Softmax\n",
     "\n",
     "Reference solution for the numerically-stable Softmax function.\n",
@@ -25,6 +26,20 @@
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -73,4 +88,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/03_linear_solution.ipynb

@@ -17,12 +17,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/03_linear_solution.ipynb)\n\n",
     "# 🟡 Solution: Simple Linear Layer\n",
     "\n",
     "Reference solution for a fully-connected linear layer: **y = xW^T + b**"
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -77,4 +92,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/04_layernorm_solution.ipynb

@@ -17,6 +17,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/04_layernorm_solution.ipynb)\n\n",
     "# 🟡 Solution: Implement LayerNorm\n",
     "\n",
     "Reference solution for Layer Normalization.\n",
@@ -25,6 +26,20 @@
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -76,4 +91,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/05_attention_solution.ipynb

@@ -5,6 +5,7 @@
    "id": "5f63d076",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/05_attention_solution.ipynb)\n\n",
     "# 🔴 Solution: Softmax Attention\n",
     "\n",
     "Reference solution for the core Transformer attention mechanism.\n",
@@ -12,6 +13,21 @@
     "$$\\text{Attention}(Q, K, V) = \\text{softmax}\\!\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V$$"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ce663fb0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

solutions/06_multihead_attention_solution.ipynb

@@ -1,105 +1,120 @@
 {
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "# 🔴 Solution: Multi-Head Attention\n",
-        "\n",
-        "Reference solution for the Multi-Head Attention mechanism.\n",
-        "\n",
-        "$$\\text{MultiHead}(Q, K, V) = \\text{Concat}(\\text{head}_1, \\dots, \\text{head}_h) W^O$$"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import torch\n",
-        "import torch.nn as nn\n",
-        "import math"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "46b73737",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# ✅ SOLUTION\n",
-        "\n",
-        "class MultiHeadAttention:\n",
-        "    def __init__(self, d_model: int, num_heads: int):\n",
-        "        self.num_heads = num_heads\n",
-        "        self.d_k = d_model // num_heads\n",
-        "\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",
-        "\n",
-        "    def forward(self, Q, K, V):\n",
-        "        B, S_q, _ = Q.shape\n",
-        "        S_k = K.shape[1]\n",
-        "\n",
-        "        q = self.W_q(Q).view(B, S_q, self.num_heads, self.d_k).transpose(1, 2)\n",
-        "        k = self.W_k(K).view(B, S_k, self.num_heads, self.d_k).transpose(1, 2)\n",
-        "        v = self.W_v(V).view(B, S_k, self.num_heads, self.d_k).transpose(1, 2)\n",
-        "\n",
-        "        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)\n",
-        "        weights = torch.softmax(scores, dim=-1)\n",
-        "        attn = torch.matmul(weights, v)\n",
-        "\n",
-        "        out = attn.transpose(1, 2).contiguous().view(B, S_q, -1)\n",
-        "        return self.W_o(out)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Verify\n",
-        "torch.manual_seed(0)\n",
-        "mha = MultiHeadAttention(d_model=32, num_heads=4)\n",
-        "x = torch.randn(2, 6, 32)\n",
-        "out = mha.forward(x, x, x)\n",
-        "print(\"Self-attn shape:\", out.shape)\n",
-        "\n",
-        "Q = torch.randn(1, 3, 32)\n",
-        "K = torch.randn(1, 7, 32)\n",
-        "V = torch.randn(1, 7, 32)\n",
-        "out2 = mha.forward(Q, K, V)\n",
-        "print(\"Cross-attn shape:\", out2.shape)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Run judge\n",
-        "from torch_judge import check\n",
-        "check(\"mha\")"
-      ]
-    }
-  ],
-  "metadata": {
-    "kernelspec": {
-      "display_name": "Python 3",
-      "language": "python",
-      "name": "python3"
-    },
-    "language_info": {
-      "name": "python",
-      "version": "3.11.0"
-    }
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/06_multihead_attention_solution.ipynb)\n\n",
+    "# 🔴 Solution: Multi-Head Attention\n",
+    "\n",
+    "Reference solution for the Multi-Head Attention mechanism.\n",
+    "\n",
+    "$$\\text{MultiHead}(Q, K, V) = \\text{Concat}(\\text{head}_1, \\dots, \\text{head}_h) W^O$$"
+   ]
   },
-  "nbformat": 4,
-  "nbformat_minor": 5
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import math"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "46b73737",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ✅ SOLUTION\n",
+    "\n",
+    "class MultiHeadAttention:\n",
+    "    def __init__(self, d_model: int, num_heads: int):\n",
+    "        self.num_heads = num_heads\n",
+    "        self.d_k = d_model // num_heads\n",
+    "\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",
+    "\n",
+    "    def forward(self, Q, K, V):\n",
+    "        B, S_q, _ = Q.shape\n",
+    "        S_k = K.shape[1]\n",
+    "\n",
+    "        q = self.W_q(Q).view(B, S_q, self.num_heads, self.d_k).transpose(1, 2)\n",
+    "        k = self.W_k(K).view(B, S_k, self.num_heads, self.d_k).transpose(1, 2)\n",
+    "        v = self.W_v(V).view(B, S_k, self.num_heads, self.d_k).transpose(1, 2)\n",
+    "\n",
+    "        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)\n",
+    "        weights = torch.softmax(scores, dim=-1)\n",
+    "        attn = torch.matmul(weights, v)\n",
+    "\n",
+    "        out = attn.transpose(1, 2).contiguous().view(B, S_q, -1)\n",
+    "        return self.W_o(out)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Verify\n",
+    "torch.manual_seed(0)\n",
+    "mha = MultiHeadAttention(d_model=32, num_heads=4)\n",
+    "x = torch.randn(2, 6, 32)\n",
+    "out = mha.forward(x, x, x)\n",
+    "print(\"Self-attn shape:\", out.shape)\n",
+    "\n",
+    "Q = torch.randn(1, 3, 32)\n",
+    "K = torch.randn(1, 7, 32)\n",
+    "V = torch.randn(1, 7, 32)\n",
+    "out2 = mha.forward(Q, K, V)\n",
+    "print(\"Cross-attn shape:\", out2.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Run judge\n",
+    "from torch_judge import check\n",
+    "check(\"mha\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
 }

solutions/07_batchnorm_solution.ipynb

@@ -1,117 +1,132 @@
 {
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "id": "ffd42526",
-      "metadata": {},
-      "source": [
-        "# 🟡 Solution: Implement BatchNorm\n",
-        "\n",
-        "Reference solution for Batch Normalization with both **training** and **inference** behavior, including running mean/variance updates."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import torch"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "70488b9f",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# ✅ SOLUTION\n",
-        "\n",
-        "import torch\n",
-        "\n",
-        "def my_batch_norm(\n",
-        "    x,\n",
-        "    gamma,\n",
-        "    beta,\n",
-        "    running_mean,\n",
-        "    running_var,\n",
-        "    eps=1e-5,\n",
-        "    momentum=0.1,\n",
-        "    training=True,\n",
-        "):\n",
-        "    \"\"\"BatchNorm with train/eval behavior and running stats.\n",
-        "\n",
-        "    - Training: use batch stats, update running_mean / running_var in-place.\n",
-        "    - Inference: use running_mean / running_var as-is.\n",
-        "    \"\"\"\n",
-        "    if training:\n",
-        "        batch_mean = x.mean(dim=0)\n",
-        "        batch_var = x.var(dim=0, unbiased=False)\n",
-        "\n",
-        "        # Update running statistics in-place. Detach to avoid tracking gradients.\n",
-        "        running_mean.mul_(1 - momentum).add_(momentum * batch_mean.detach())\n",
-        "        running_var.mul_(1 - momentum).add_(momentum * batch_var.detach())\n",
-        "\n",
-        "        mean = batch_mean\n",
-        "        var = batch_var\n",
-        "    else:\n",
-        "        mean = running_mean\n",
-        "        var = running_var\n",
-        "\n",
-        "    x_norm = (x - mean) / torch.sqrt(var + eps)\n",
-        "    return gamma * x_norm + beta"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "dbd7bb4e",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Verify\n",
-        "x = torch.randn(8, 4)\n",
-        "gamma = torch.ones(4)\n",
-        "beta = torch.zeros(4)\n",
-        "\n",
-        "running_mean = torch.zeros(4)\n",
-        "running_var = torch.ones(4)\n",
-        "\n",
-        "# Training behavior: normalize with batch stats and update running stats\n",
-        "out_train = my_batch_norm(x, gamma, beta, running_mean, running_var, training=True)\n",
-        "print(\"[Train] Column means:\", out_train.mean(dim=0))\n",
-        "print(\"[Train] Column stds: \", out_train.std(dim=0))\n",
-        "print(\"Updated running_mean:\", running_mean)\n",
-        "print(\"Updated running_var:\", running_var)\n",
-        "\n",
-        "# Inference behavior: use running_mean / running_var only\n",
-        "out_eval = my_batch_norm(x, gamma, beta, running_mean, running_var, training=False)\n",
-        "print(\"[Eval] Column means (using running stats):\", out_eval.mean(dim=0))"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "from torch_judge import check\n",
-        "check('batchnorm')"
-      ]
-    }
-  ],
-  "metadata": {
-    "kernelspec": {
-      "display_name": "Python 3",
-      "language": "python",
-      "name": "python3"
-    },
-    "language_info": {
-      "name": "python",
-      "version": "3.11.0"
-    }
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "ffd42526",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/07_batchnorm_solution.ipynb)\n\n",
+    "# 🟡 Solution: Implement BatchNorm\n",
+    "\n",
+    "Reference solution for Batch Normalization with both **training** and **inference** behavior, including running mean/variance updates."
+   ]
   },
-  "nbformat": 4,
-  "nbformat_minor": 5
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "70488b9f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ✅ SOLUTION\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "def my_batch_norm(\n",
+    "    x,\n",
+    "    gamma,\n",
+    "    beta,\n",
+    "    running_mean,\n",
+    "    running_var,\n",
+    "    eps=1e-5,\n",
+    "    momentum=0.1,\n",
+    "    training=True,\n",
+    "):\n",
+    "    \"\"\"BatchNorm with train/eval behavior and running stats.\n",
+    "\n",
+    "    - Training: use batch stats, update running_mean / running_var in-place.\n",
+    "    - Inference: use running_mean / running_var as-is.\n",
+    "    \"\"\"\n",
+    "    if training:\n",
+    "        batch_mean = x.mean(dim=0)\n",
+    "        batch_var = x.var(dim=0, unbiased=False)\n",
+    "\n",
+    "        # Update running statistics in-place. Detach to avoid tracking gradients.\n",
+    "        running_mean.mul_(1 - momentum).add_(momentum * batch_mean.detach())\n",
+    "        running_var.mul_(1 - momentum).add_(momentum * batch_var.detach())\n",
+    "\n",
+    "        mean = batch_mean\n",
+    "        var = batch_var\n",
+    "    else:\n",
+    "        mean = running_mean\n",
+    "        var = running_var\n",
+    "\n",
+    "    x_norm = (x - mean) / torch.sqrt(var + eps)\n",
+    "    return gamma * x_norm + beta"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dbd7bb4e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Verify\n",
+    "x = torch.randn(8, 4)\n",
+    "gamma = torch.ones(4)\n",
+    "beta = torch.zeros(4)\n",
+    "\n",
+    "running_mean = torch.zeros(4)\n",
+    "running_var = torch.ones(4)\n",
+    "\n",
+    "# Training behavior: normalize with batch stats and update running stats\n",
+    "out_train = my_batch_norm(x, gamma, beta, running_mean, running_var, training=True)\n",
+    "print(\"[Train] Column means:\", out_train.mean(dim=0))\n",
+    "print(\"[Train] Column stds: \", out_train.std(dim=0))\n",
+    "print(\"Updated running_mean:\", running_mean)\n",
+    "print(\"Updated running_var:\", running_var)\n",
+    "\n",
+    "# Inference behavior: use running_mean / running_var only\n",
+    "out_eval = my_batch_norm(x, gamma, beta, running_mean, running_var, training=False)\n",
+    "print(\"[Eval] Column means (using running stats):\", out_eval.mean(dim=0))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torch_judge import check\n",
+    "check('batchnorm')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
 }

solutions/08_rmsnorm_solution.ipynb

@@ -17,12 +17,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/08_rmsnorm_solution.ipynb)\n\n",
     "# 🟡 Solution: Implement RMSNorm\n",
     "\n",
     "Reference solution for Root Mean Square Normalization."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -67,4 +82,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/09_causal_attention_solution.ipynb

@@ -17,12 +17,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/09_causal_attention_solution.ipynb)\n\n",
     "# 🔴 Solution: Causal Self-Attention\n",
     "\n",
     "Reference solution — softmax attention with an upper-triangular mask."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -77,4 +92,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/10_gqa_solution.ipynb

@@ -17,12 +17,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/10_gqa_solution.ipynb)\n\n",
     "# 🔴 Solution: Grouped Query Attention\n",
     "\n",
     "Reference solution for GQA — MHA with shared KV heads."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -88,4 +103,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/11_sliding_window_solution.ipynb

@@ -17,12 +17,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/11_sliding_window_solution.ipynb)\n\n",
     "# 🔴 Solution: Sliding Window Attention\n",
     "\n",
     "Reference solution — softmax attention with a band mask."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -73,4 +88,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/12_linear_attention_solution.ipynb

@@ -17,12 +17,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/12_linear_attention_solution.ipynb)\n\n",
     "# 🔴 Solution: Linear Self-Attention\n",
     "\n",
     "Reference solution — kernel-based attention with elu+1 feature map."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -72,4 +87,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/13_gpt2_block_solution.ipynb

@@ -17,12 +17,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/13_gpt2_block_solution.ipynb)\n\n",
     "# 🔴 Solution: GPT-2 Transformer Block\n",
     "\n",
     "Reference solution — pre-norm, causal self-attention, 4x MLP with GELU."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -102,4 +117,4 @@
    "execution_count": null
   }
  ]
-}
+}

solutions/14_kv_cache_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# \ud83d\udd34 Solution: KV Cache Attention\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/14_kv_cache_solution.ipynb)\n\n",
+    "# 🔴 Solution: KV Cache Attention\n",
     "\n",
-    "Reference solution \u2014 multi-head attention with KV caching for autoregressive inference."
+    "Reference solution — multi-head attention with KV caching for autoregressive inference."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -26,7 +41,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# \u2705 SOLUTION\n",
+    "# ✅ SOLUTION\n",
     "\n",
     "class KVCacheAttention(nn.Module):\n",
     "    def __init__(self, d_model, num_heads):\n",

solutions/15_mlp_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# \ud83d\udfe0 Solution: SwiGLU MLP\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/15_mlp_solution.ipynb)\n\n",
+    "# 🟠 Solution: SwiGLU MLP\n",
     "\n",
-    "Reference solution \u2014 gated feed-forward network used in LLaMA, Mistral, and PaLM."
+    "Reference solution — gated feed-forward network used in LLaMA, Mistral, and PaLM."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -26,7 +41,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# \u2705 SOLUTION\n",
+    "# ✅ SOLUTION\n",
     "\n",
     "class SwiGLUMLP(nn.Module):\n",
     "    def __init__(self, d_model, d_ff):\n",

solutions/16_cross_entropy_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/16_cross_entropy_solution.ipynb)\n\n",
     "# Solution: Cross-Entropy Loss\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -69,4 +84,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/17_dropout_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/17_dropout_solution.ipynb)\n\n",
     "# Solution: Implement Dropout\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -79,4 +94,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/18_embedding_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/18_embedding_solution.ipynb)\n\n",
     "# Solution: Embedding Layer\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -74,4 +89,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/19_gelu_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/19_gelu_solution.ipynb)\n\n",
     "# Solution: GELU Activation\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -68,4 +83,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/20_weight_init_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/20_weight_init_solution.ipynb)\n\n",
     "# Solution: Kaiming Initialization\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -74,4 +89,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/21_gradient_clipping_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/21_gradient_clipping_solution.ipynb)\n\n",
     "# Solution: Gradient Norm Clipping\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -76,4 +91,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/22_conv2d_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/22_conv2d_solution.ipynb)\n\n",
     "# Solution: 2D Convolution\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -79,4 +94,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/23_cross_attention_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/23_cross_attention_solution.ipynb)\n\n",
     "# Solution: Multi-Head Cross-Attention\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -88,4 +103,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/24_rope_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/24_rope_solution.ipynb)\n\n",
     "# Solution: Rotary Position Embedding (RoPE)\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -83,4 +98,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/25_flash_attention_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/25_flash_attention_solution.ipynb)\n\n",
     "# Solution: Flash Attention (Tiled)\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -92,4 +107,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/26_lora_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/26_lora_solution.ipynb)\n\n",
     "# Solution: LoRA (Low-Rank Adaptation)\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -81,4 +96,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/27_vit_patch_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/27_vit_patch_solution.ipynb)\n\n",
     "# Solution: ViT Patch Embedding\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -81,4 +96,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/28_moe_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/28_moe_solution.ipynb)\n\n",
     "# Solution: Mixture of Experts (MoE)\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -94,4 +109,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/29_adam_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/29_adam_solution.ipynb)\n\n",
     "# Solution: Adam Optimizer\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -97,4 +112,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/30_cosine_lr_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/30_cosine_lr_solution.ipynb)\n\n",
     "# Solution: Cosine LR Scheduler with Warmup\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -71,4 +86,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/31_gradient_accumulation_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/31_gradient_accumulation_solution.ipynb)\n\n",
     "# Solution: Gradient Accumulation\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -78,4 +93,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/32_topk_sampling_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/32_topk_sampling_solution.ipynb)\n\n",
     "# Solution: Top-k / Top-p Sampling\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -79,4 +94,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/33_beam_search_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/33_beam_search_solution.ipynb)\n\n",
     "# Solution: Beam Search Decoding\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -88,4 +103,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/34_speculative_decoding_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/34_speculative_decoding_solution.ipynb)\n\n",
     "# Solution: Speculative Decoding\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -84,4 +99,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/35_bpe_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/35_bpe_solution.ipynb)\n\n",
     "# Solution: Byte-Pair Encoding (BPE)\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -113,4 +128,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/36_int8_quantization_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/36_int8_quantization_solution.ipynb)\n\n",
     "# Solution: INT8 Quantized Linear\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -83,4 +98,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/37_dpo_loss_solution.ipynb

@@ -4,12 +4,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/37_dpo_loss_solution.ipynb)\n\n",
     "# Solution: DPO (Direct Preference Optimization) Loss\n",
     "\n",
     "Reference solution."
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},
@@ -73,4 +88,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

solutions/38_grpo_loss_solution.ipynb

@@ -4,11 +4,26 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/38_grpo_loss_solution.ipynb)\n\n",
     "# Solution: GRPO (Group Relative Policy Optimization) Loss\n",
     "\n",
     "Reference solution."
    ]
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "execution_count": null,

solutions/39_ppo_loss_solution.ipynb

@@ -4,11 +4,26 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/39_ppo_loss_solution.ipynb)\n\n",
     "# Solution: PPO Clipped Loss\n",
     "\n",
     "Reference solution for the PPO clipped surrogate loss task.\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -89,4 +104,3 @@
  "nbformat": 4,
  "nbformat_minor": 5
 }
-

solutions/40_linear_regression_solution.ipynb

@@ -1,125 +1,140 @@
 {
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "# 🟡 Solution: Linear Regression\n",
-        "\n",
-        "Reference solution demonstrating closed-form, gradient descent, and nn.Linear approaches."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import torch\n",
-        "import torch.nn as nn"
-      ],
-      "execution_count": null
-    },
-    {
-      "cell_type": "code",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# ✅ SOLUTION\n",
-        "\n",
-        "class LinearRegression:\n",
-        "    def closed_form(self, X: torch.Tensor, y: torch.Tensor):\n",
-        "        \"\"\"Normal equation via augmented matrix.\"\"\"\n",
-        "        N, D = X.shape\n",
-        "        # Augment X with ones column for bias\n",
-        "        X_aug = torch.cat([X, torch.ones(N, 1)], dim=1)  # (N, D+1)\n",
-        "        # Solve (X^T X) theta = X^T y\n",
-        "        theta = torch.linalg.lstsq(X_aug, y).solution      # (D+1,)\n",
-        "        w = theta[:D]\n",
-        "        b = theta[D]\n",
-        "        return w.detach(), b.detach()\n",
-        "\n",
-        "    def gradient_descent(self, X: torch.Tensor, y: torch.Tensor,\n",
-        "                         lr: float = 0.01, steps: int = 1000):\n",
-        "        \"\"\"Manual gradient computation — no autograd.\"\"\"\n",
-        "        N, D = X.shape\n",
-        "        w = torch.zeros(D)\n",
-        "        b = torch.tensor(0.0)\n",
-        "\n",
-        "        for _ in range(steps):\n",
-        "            pred = X @ w + b          # (N,)\n",
-        "            error = pred - y           # (N,)\n",
-        "            grad_w = (2.0 / N) * (X.T @ error)  # (D,)\n",
-        "            grad_b = (2.0 / N) * error.sum()     # scalar\n",
-        "            w = w - lr * grad_w\n",
-        "            b = b - lr * grad_b\n",
-        "\n",
-        "        return w, b\n",
-        "\n",
-        "    def nn_linear(self, X: torch.Tensor, y: torch.Tensor,\n",
-        "                  lr: float = 0.01, steps: int = 1000):\n",
-        "        \"\"\"PyTorch nn.Linear with autograd training loop.\"\"\"\n",
-        "        N, D = X.shape\n",
-        "        layer = nn.Linear(D, 1)\n",
-        "        optimizer = torch.optim.SGD(layer.parameters(), lr=lr)\n",
-        "        loss_fn = nn.MSELoss()\n",
-        "\n",
-        "        for _ in range(steps):\n",
-        "            optimizer.zero_grad()\n",
-        "            pred = layer(X).squeeze(-1)  # (N,)\n",
-        "            loss = loss_fn(pred, y)\n",
-        "            loss.backward()\n",
-        "            optimizer.step()\n",
-        "\n",
-        "        w = layer.weight.data.squeeze(0)  # (D,)\n",
-        "        b = layer.bias.data.squeeze(0)    # scalar ()\n",
-        "        return w, b"
-      ],
-      "execution_count": null
-    },
-    {
-      "cell_type": "code",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Verify\n",
-        "torch.manual_seed(42)\n",
-        "X = torch.randn(100, 3)\n",
-        "true_w = torch.tensor([2.0, -1.0, 0.5])\n",
-        "y = X @ true_w + 3.0\n",
-        "\n",
-        "model = LinearRegression()\n",
-        "for name, method in [(\"Closed-form\", model.closed_form),\n",
-        "                      (\"Grad Descent\", lambda X, y: model.gradient_descent(X, y, lr=0.05, steps=2000)),\n",
-        "                      (\"nn.Linear\", lambda X, y: model.nn_linear(X, y, lr=0.05, steps=2000))]:\n",
-        "    w, b = method(X, y)\n",
-        "    print(f\"{name:13s}  w={w.tolist()}  b={b.item():.4f}\")\n",
-        "print(f\"{'True':13s}  w={true_w.tolist()}  b=3.0000\")"
-      ],
-      "execution_count": null
-    },
-    {
-      "cell_type": "code",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# ✅ SUBMIT\n",
-        "from torch_judge import check\n",
-        "check(\"linear_regression\")"
-      ],
-      "execution_count": null
-    }
-  ],
-  "metadata": {
-    "kernelspec": {
-      "display_name": "Python 3",
-      "language": "python",
-      "name": "python3"
-    },
-    "language_info": {
-      "name": "python",
-      "version": "3.11.0"
-    }
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/40_linear_regression_solution.ipynb)\n\n",
+    "# 🟡 Solution: Linear Regression\n",
+    "\n",
+    "Reference solution demonstrating closed-form, gradient descent, and nn.Linear approaches."
+   ]
   },
-  "nbformat": 4,
-  "nbformat_minor": 4
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn"
+   ],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ✅ SOLUTION\n",
+    "\n",
+    "class LinearRegression:\n",
+    "    def closed_form(self, X: torch.Tensor, y: torch.Tensor):\n",
+    "        \"\"\"Normal equation via augmented matrix.\"\"\"\n",
+    "        N, D = X.shape\n",
+    "        # Augment X with ones column for bias\n",
+    "        X_aug = torch.cat([X, torch.ones(N, 1)], dim=1)  # (N, D+1)\n",
+    "        # Solve (X^T X) theta = X^T y\n",
+    "        theta = torch.linalg.lstsq(X_aug, y).solution      # (D+1,)\n",
+    "        w = theta[:D]\n",
+    "        b = theta[D]\n",
+    "        return w.detach(), b.detach()\n",
+    "\n",
+    "    def gradient_descent(self, X: torch.Tensor, y: torch.Tensor,\n",
+    "                         lr: float = 0.01, steps: int = 1000):\n",
+    "        \"\"\"Manual gradient computation — no autograd.\"\"\"\n",
+    "        N, D = X.shape\n",
+    "        w = torch.zeros(D)\n",
+    "        b = torch.tensor(0.0)\n",
+    "\n",
+    "        for _ in range(steps):\n",
+    "            pred = X @ w + b          # (N,)\n",
+    "            error = pred - y           # (N,)\n",
+    "            grad_w = (2.0 / N) * (X.T @ error)  # (D,)\n",
+    "            grad_b = (2.0 / N) * error.sum()     # scalar\n",
+    "            w = w - lr * grad_w\n",
+    "            b = b - lr * grad_b\n",
+    "\n",
+    "        return w, b\n",
+    "\n",
+    "    def nn_linear(self, X: torch.Tensor, y: torch.Tensor,\n",
+    "                  lr: float = 0.01, steps: int = 1000):\n",
+    "        \"\"\"PyTorch nn.Linear with autograd training loop.\"\"\"\n",
+    "        N, D = X.shape\n",
+    "        layer = nn.Linear(D, 1)\n",
+    "        optimizer = torch.optim.SGD(layer.parameters(), lr=lr)\n",
+    "        loss_fn = nn.MSELoss()\n",
+    "\n",
+    "        for _ in range(steps):\n",
+    "            optimizer.zero_grad()\n",
+    "            pred = layer(X).squeeze(-1)  # (N,)\n",
+    "            loss = loss_fn(pred, y)\n",
+    "            loss.backward()\n",
+    "            optimizer.step()\n",
+    "\n",
+    "        w = layer.weight.data.squeeze(0)  # (D,)\n",
+    "        b = layer.bias.data.squeeze(0)    # scalar ()\n",
+    "        return w, b"
+   ],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Verify\n",
+    "torch.manual_seed(42)\n",
+    "X = torch.randn(100, 3)\n",
+    "true_w = torch.tensor([2.0, -1.0, 0.5])\n",
+    "y = X @ true_w + 3.0\n",
+    "\n",
+    "model = LinearRegression()\n",
+    "for name, method in [(\"Closed-form\", model.closed_form),\n",
+    "                      (\"Grad Descent\", lambda X, y: model.gradient_descent(X, y, lr=0.05, steps=2000)),\n",
+    "                      (\"nn.Linear\", lambda X, y: model.nn_linear(X, y, lr=0.05, steps=2000))]:\n",
+    "    w, b = method(X, y)\n",
+    "    print(f\"{name:13s}  w={w.tolist()}  b={b.item():.4f}\")\n",
+    "print(f\"{'True':13s}  w={true_w.tolist()}  b=3.0000\")"
+   ],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ✅ SUBMIT\n",
+    "from torch_judge import check\n",
+    "check(\"linear_regression\")"
+   ],
+   "execution_count": null
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
 }

templates/00_welcome.ipynb

@@ -31,24 +31,120 @@
     "\n",
     "> 💡 Every notebook also has a **Colab** toolbar button and an **Open in Colab** badge — use them to run problems in Google Colab with zero setup.\n",
     "\n",
-    "## Quick Start"
+    "## Quick Start\n",
+    "\n",
+    "📖 **Reference solutions in Colab**: [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/01_relu_solution.ipynb) — Start with ReLU. Or use the **Colab** links in the table below for each solution."
    ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "from torch_judge import status\n",
     "status()"
-   ],
-   "execution_count": null,
-   "outputs": []
+   ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Problem List (40 problems)\n\n### 🧱 Fundamentals — \"Implement X from scratch\"\n\n| # | Problem | Difficulty | Template | Solution |\n|:---:|---------|:----------:|:--------:|:--------:|\n| 1 | ReLU | 🟢 Easy | [Open](01_relu.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/01_relu.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/01_relu.ipynb\" target=\"_blank\">Colab</a> | [Open](01_relu_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/01_relu_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/01_relu_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 2 | Softmax | 🟢 Easy | [Open](02_softmax.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/02_softmax.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/02_softmax.ipynb\" target=\"_blank\">Colab</a> | [Open](02_softmax_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/02_softmax_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/02_softmax_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 16 | Cross-Entropy Loss | 🟢 Easy | [Open](16_cross_entropy.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/16_cross_entropy.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/16_cross_entropy.ipynb\" target=\"_blank\">Colab</a> | [Open](16_cross_entropy_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/16_cross_entropy_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/16_cross_entropy_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 17 | Dropout | 🟢 Easy | [Open](17_dropout.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/17_dropout.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/17_dropout.ipynb\" target=\"_blank\">Colab</a> | [Open](17_dropout_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/17_dropout_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/17_dropout_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 18 | Embedding | 🟢 Easy | [Open](18_embedding.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/18_embedding.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/18_embedding.ipynb\" target=\"_blank\">Colab</a> | [Open](18_embedding_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/18_embedding_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/18_embedding_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 19 | GELU | 🟢 Easy | [Open](19_gelu.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/19_gelu.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/19_gelu.ipynb\" target=\"_blank\">Colab</a> | [Open](19_gelu_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/19_gelu_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/19_gelu_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 20 | Kaiming Init | 🟢 Easy | [Open](20_weight_init.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/20_weight_init.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/20_weight_init.ipynb\" target=\"_blank\">Colab</a> | [Open](20_weight_init_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/20_weight_init_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/20_weight_init_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 21 | Gradient Clipping | 🟢 Easy | [Open](21_gradient_clipping.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/21_gradient_clipping.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/21_gradient_clipping.ipynb\" target=\"_blank\">Colab</a> | [Open](21_gradient_clipping_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/21_gradient_clipping_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/21_gradient_clipping_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 31 | Gradient Accumulation | 🟢 Easy | [Open](31_gradient_accumulation.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/31_gradient_accumulation.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/31_gradient_accumulation.ipynb\" target=\"_blank\">Colab</a> | [Open](31_gradient_accumulation_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/31_gradient_accumulation_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/31_gradient_accumulation_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 3 | Linear Layer | 🟡 Medium | [Open](03_linear.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/03_linear.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/03_linear.ipynb\" target=\"_blank\">Colab</a> | [Open](03_linear_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/03_linear_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/03_linear_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 4 | LayerNorm | 🟡 Medium | [Open](04_layernorm.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/04_layernorm.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/04_layernorm.ipynb\" target=\"_blank\">Colab</a> | [Open](04_layernorm_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/04_layernorm_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/04_layernorm_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 7 | BatchNorm | 🟡 Medium | [Open](07_batchnorm.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/07_batchnorm.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/07_batchnorm.ipynb\" target=\"_blank\">Colab</a> | [Open](07_batchnorm_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/07_batchnorm_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/07_batchnorm_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 8 | RMSNorm | 🟡 Medium | [Open](08_rmsnorm.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/08_rmsnorm.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/08_rmsnorm.ipynb\" target=\"_blank\">Colab</a> | [Open](08_rmsnorm_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/08_rmsnorm_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/08_rmsnorm_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 15 | SwiGLU MLP | 🟡 Medium | [Open](15_mlp.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/15_mlp.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/15_mlp.ipynb\" target=\"_blank\">Colab</a> | [Open](15_mlp_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/15_mlp_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/15_mlp_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 22 | Conv2d | 🟡 Medium | [Open](22_conv2d.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/22_conv2d.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/22_conv2d.ipynb\" target=\"_blank\">Colab</a> | [Open](22_conv2d_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/22_conv2d_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/22_conv2d_solution.ipynb\" target=\"_blank\">Colab</a> |\n\n### 🧠 Attention Mechanisms\n\n| # | Problem | Difficulty | Template | Solution |\n|:---:|---------|:----------:|:--------:|:--------:|\n| 23 | Cross-Attention | 🟡 Medium | [Open](23_cross_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/23_cross_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/23_cross_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](23_cross_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/23_cross_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/23_cross_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 5 | Scaled Dot-Product Attention | 🔴 Hard | [Open](05_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/05_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/05_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](05_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/05_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/05_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 6 | Multi-Head Attention | 🔴 Hard | [Open](06_multihead_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/06_multihead_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/06_multihead_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](06_multihead_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/06_multihead_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/06_multihead_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 9 | Causal Self-Attention | 🔴 Hard | [Open](09_causal_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/09_causal_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/09_causal_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](09_causal_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/09_causal_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/09_causal_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 10 | Grouped Query Attention | 🔴 Hard | [Open](10_gqa.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/10_gqa.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/10_gqa.ipynb\" target=\"_blank\">Colab</a> | [Open](10_gqa_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/10_gqa_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/10_gqa_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 11 | Sliding Window Attention | 🔴 Hard | [Open](11_sliding_window.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/11_sliding_window.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/11_sliding_window.ipynb\" target=\"_blank\">Colab</a> | [Open](11_sliding_window_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/11_sliding_window_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/11_sliding_window_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 12 | Linear Attention | 🔴 Hard | [Open](12_linear_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/12_linear_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/12_linear_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](12_linear_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/12_linear_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/12_linear_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 14 | KV Cache Attention | 🔴 Hard | [Open](14_kv_cache.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/14_kv_cache.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/14_kv_cache.ipynb\" target=\"_blank\">Colab</a> | [Open](14_kv_cache_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/14_kv_cache_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/14_kv_cache_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 24 | RoPE | 🔴 Hard | [Open](24_rope.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/24_rope.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/24_rope.ipynb\" target=\"_blank\">Colab</a> | [Open](24_rope_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/24_rope_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/24_rope_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 25 | Flash Attention | 🔴 Hard | [Open](25_flash_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/25_flash_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/25_flash_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](25_flash_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/25_flash_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/25_flash_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n\n### 🏗️ Architecture & Adaptation\n\n| # | Problem | Difficulty | Template | Solution |\n|:---:|---------|:----------:|:--------:|:--------:|\n| 26 | LoRA | 🟡 Medium | [Open](26_lora.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/26_lora.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/26_lora.ipynb\" target=\"_blank\">Colab</a> | [Open](26_lora_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/26_lora_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/26_lora_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 27 | ViT Patch Embedding | 🟡 Medium | [Open](27_vit_patch.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/27_vit_patch.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/27_vit_patch.ipynb\" target=\"_blank\">Colab</a> | [Open](27_vit_patch_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/27_vit_patch_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/27_vit_patch_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 13 | GPT-2 Block | 🔴 Hard | [Open](13_gpt2_block.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/13_gpt2_block.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/13_gpt2_block.ipynb\" target=\"_blank\">Colab</a> | [Open](13_gpt2_block_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/13_gpt2_block_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/13_gpt2_block_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 28 | Mixture of Experts | 🔴 Hard | [Open](28_moe.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/28_moe.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/28_moe.ipynb\" target=\"_blank\">Colab</a> | [Open](28_moe_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/28_moe_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/28_moe_solution.ipynb\" target=\"_blank\">Colab</a> |\n\n### ⚙️ Training & Optimization\n\n| # | Problem | Difficulty | Template | Solution |\n|:---:|---------|:----------:|:--------:|:--------:|\n| 29 | Adam Optimizer | 🟡 Medium | [Open](29_adam.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/29_adam.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/29_adam.ipynb\" target=\"_blank\">Colab</a> | [Open](29_adam_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/29_adam_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/29_adam_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 30 | Cosine LR Scheduler | 🟡 Medium | [Open](30_cosine_lr.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/30_cosine_lr.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/30_cosine_lr.ipynb\" target=\"_blank\">Colab</a> | [Open](30_cosine_lr_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/30_cosine_lr_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/30_cosine_lr_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 40 | Linear Regression | 🟡 Medium | [Open](40_linear_regression.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/40_linear_regression.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/40_linear_regression.ipynb\" target=\"_blank\">Colab</a> | [Open](40_linear_regression_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/40_linear_regression_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/40_linear_regression_solution.ipynb\" target=\"_blank\">Colab</a> |\n\n### 🎯 Inference & Decoding\n\n| # | Problem | Difficulty | Template | Solution |\n|:---:|---------|:----------:|:--------:|:--------:|\n| 32 | Top-k / Top-p Sampling | 🟡 Medium | [Open](32_topk_sampling.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/32_topk_sampling.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/32_topk_sampling.ipynb\" target=\"_blank\">Colab</a> | [Open](32_topk_sampling_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/32_topk_sampling_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/32_topk_sampling_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 33 | Beam Search | 🟡 Medium | [Open](33_beam_search.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/33_beam_search.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/33_beam_search.ipynb\" target=\"_blank\">Colab</a> | [Open](33_beam_search_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/33_beam_search_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/33_beam_search_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 34 | Speculative Decoding | 🔴 Hard | [Open](34_speculative_decoding.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/34_speculative_decoding.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/34_speculative_decoding.ipynb\" target=\"_blank\">Colab</a> | [Open](34_speculative_decoding_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/34_speculative_decoding_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/34_speculative_decoding_solution.ipynb\" target=\"_blank\">Colab</a> |\n\n### 🔬 Advanced\n\n| # | Problem | Difficulty | Template | Solution |\n|:---:|---------|:----------:|:--------:|:--------:|\n| 35 | BPE Tokenizer | 🔴 Hard | [Open](35_bpe.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/35_bpe.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/35_bpe.ipynb\" target=\"_blank\">Colab</a> | [Open](35_bpe_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/35_bpe_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/35_bpe_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 36 | INT8 Quantization | 🔴 Hard | [Open](36_int8_quantization.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/36_int8_quantization.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/36_int8_quantization.ipynb\" target=\"_blank\">Colab</a> | [Open](36_int8_quantization_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/36_int8_quantization_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/36_int8_quantization_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 37 | DPO Loss | 🔴 Hard | [Open](37_dpo_loss.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/37_dpo_loss.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/37_dpo_loss.ipynb\" target=\"_blank\">Colab</a> | [Open](37_dpo_loss_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/37_dpo_loss_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/37_dpo_loss_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 38 | GRPO Loss | 🔴 Hard | [Open](38_grpo_loss.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/38_grpo_loss.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/38_grpo_loss.ipynb\" target=\"_blank\">Colab</a> | [Open](38_grpo_loss_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/38_grpo_loss_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/38_grpo_loss_solution.ipynb\" target=\"_blank\">Colab</a> |\n| 39 | PPO Loss | 🔴 Hard | [Open](39_ppo_loss.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/39_ppo_loss.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/39_ppo_loss.ipynb\" target=\"_blank\">Colab</a> | [Open](39_ppo_loss_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/39_ppo_loss_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/39_ppo_loss_solution.ipynb\" target=\"_blank\">Colab</a> |\n\n## Useful Commands\n\n```python\nfrom torch_judge import check, hint, status\n\nstatus()                   # Progress dashboard\ncheck(\"relu\")              # Judge your implementation\nhint(\"causal_attention\")   # Get a hint\n```"
+    "## Problem List (40 problems)\n",
+    "\n",
+    "### 🧱 Fundamentals — \"Implement X from scratch\"\n",
+    "\n",
+    "| # | Problem | Difficulty | Template | Solution |\n",
+    "|:---:|---------|:----------:|:--------:|:--------:|\n",
+    "| 1 | ReLU | 🟢 Easy | [Open](01_relu.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/01_relu.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/01_relu.ipynb\" target=\"_blank\">Colab</a> | [Open](01_relu_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/01_relu_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/01_relu_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 2 | Softmax | 🟢 Easy | [Open](02_softmax.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/02_softmax.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/02_softmax.ipynb\" target=\"_blank\">Colab</a> | [Open](02_softmax_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/02_softmax_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/02_softmax_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 16 | Cross-Entropy Loss | 🟢 Easy | [Open](16_cross_entropy.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/16_cross_entropy.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/16_cross_entropy.ipynb\" target=\"_blank\">Colab</a> | [Open](16_cross_entropy_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/16_cross_entropy_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/16_cross_entropy_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 17 | Dropout | 🟢 Easy | [Open](17_dropout.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/17_dropout.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/17_dropout.ipynb\" target=\"_blank\">Colab</a> | [Open](17_dropout_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/17_dropout_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/17_dropout_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 18 | Embedding | 🟢 Easy | [Open](18_embedding.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/18_embedding.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/18_embedding.ipynb\" target=\"_blank\">Colab</a> | [Open](18_embedding_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/18_embedding_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/18_embedding_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 19 | GELU | 🟢 Easy | [Open](19_gelu.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/19_gelu.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/19_gelu.ipynb\" target=\"_blank\">Colab</a> | [Open](19_gelu_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/19_gelu_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/19_gelu_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 20 | Kaiming Init | 🟢 Easy | [Open](20_weight_init.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/20_weight_init.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/20_weight_init.ipynb\" target=\"_blank\">Colab</a> | [Open](20_weight_init_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/20_weight_init_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/20_weight_init_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 21 | Gradient Clipping | 🟢 Easy | [Open](21_gradient_clipping.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/21_gradient_clipping.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/21_gradient_clipping.ipynb\" target=\"_blank\">Colab</a> | [Open](21_gradient_clipping_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/21_gradient_clipping_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/21_gradient_clipping_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 31 | Gradient Accumulation | 🟢 Easy | [Open](31_gradient_accumulation.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/31_gradient_accumulation.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/31_gradient_accumulation.ipynb\" target=\"_blank\">Colab</a> | [Open](31_gradient_accumulation_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/31_gradient_accumulation_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/31_gradient_accumulation_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 3 | Linear Layer | 🟡 Medium | [Open](03_linear.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/03_linear.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/03_linear.ipynb\" target=\"_blank\">Colab</a> | [Open](03_linear_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/03_linear_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/03_linear_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 4 | LayerNorm | 🟡 Medium | [Open](04_layernorm.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/04_layernorm.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/04_layernorm.ipynb\" target=\"_blank\">Colab</a> | [Open](04_layernorm_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/04_layernorm_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/04_layernorm_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 7 | BatchNorm | 🟡 Medium | [Open](07_batchnorm.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/07_batchnorm.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/07_batchnorm.ipynb\" target=\"_blank\">Colab</a> | [Open](07_batchnorm_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/07_batchnorm_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/07_batchnorm_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 8 | RMSNorm | 🟡 Medium | [Open](08_rmsnorm.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/08_rmsnorm.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/08_rmsnorm.ipynb\" target=\"_blank\">Colab</a> | [Open](08_rmsnorm_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/08_rmsnorm_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/08_rmsnorm_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 15 | SwiGLU MLP | 🟡 Medium | [Open](15_mlp.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/15_mlp.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/15_mlp.ipynb\" target=\"_blank\">Colab</a> | [Open](15_mlp_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/15_mlp_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/15_mlp_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 22 | Conv2d | 🟡 Medium | [Open](22_conv2d.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/22_conv2d.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/22_conv2d.ipynb\" target=\"_blank\">Colab</a> | [Open](22_conv2d_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/22_conv2d_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/22_conv2d_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "\n",
+    "### 🧠 Attention Mechanisms\n",
+    "\n",
+    "| # | Problem | Difficulty | Template | Solution |\n",
+    "|:---:|---------|:----------:|:--------:|:--------:|\n",
+    "| 23 | Cross-Attention | 🟡 Medium | [Open](23_cross_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/23_cross_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/23_cross_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](23_cross_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/23_cross_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/23_cross_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 5 | Scaled Dot-Product Attention | 🔴 Hard | [Open](05_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/05_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/05_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](05_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/05_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/05_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 6 | Multi-Head Attention | 🔴 Hard | [Open](06_multihead_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/06_multihead_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/06_multihead_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](06_multihead_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/06_multihead_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/06_multihead_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 9 | Causal Self-Attention | 🔴 Hard | [Open](09_causal_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/09_causal_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/09_causal_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](09_causal_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/09_causal_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/09_causal_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 10 | Grouped Query Attention | 🔴 Hard | [Open](10_gqa.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/10_gqa.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/10_gqa.ipynb\" target=\"_blank\">Colab</a> | [Open](10_gqa_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/10_gqa_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/10_gqa_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 11 | Sliding Window Attention | 🔴 Hard | [Open](11_sliding_window.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/11_sliding_window.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/11_sliding_window.ipynb\" target=\"_blank\">Colab</a> | [Open](11_sliding_window_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/11_sliding_window_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/11_sliding_window_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 12 | Linear Attention | 🔴 Hard | [Open](12_linear_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/12_linear_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/12_linear_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](12_linear_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/12_linear_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/12_linear_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 14 | KV Cache Attention | 🔴 Hard | [Open](14_kv_cache.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/14_kv_cache.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/14_kv_cache.ipynb\" target=\"_blank\">Colab</a> | [Open](14_kv_cache_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/14_kv_cache_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/14_kv_cache_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 24 | RoPE | 🔴 Hard | [Open](24_rope.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/24_rope.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/24_rope.ipynb\" target=\"_blank\">Colab</a> | [Open](24_rope_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/24_rope_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/24_rope_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 25 | Flash Attention | 🔴 Hard | [Open](25_flash_attention.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/25_flash_attention.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/25_flash_attention.ipynb\" target=\"_blank\">Colab</a> | [Open](25_flash_attention_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/25_flash_attention_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/25_flash_attention_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "\n",
+    "### 🏗️ Architecture & Adaptation\n",
+    "\n",
+    "| # | Problem | Difficulty | Template | Solution |\n",
+    "|:---:|---------|:----------:|:--------:|:--------:|\n",
+    "| 26 | LoRA | 🟡 Medium | [Open](26_lora.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/26_lora.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/26_lora.ipynb\" target=\"_blank\">Colab</a> | [Open](26_lora_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/26_lora_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/26_lora_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 27 | ViT Patch Embedding | 🟡 Medium | [Open](27_vit_patch.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/27_vit_patch.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/27_vit_patch.ipynb\" target=\"_blank\">Colab</a> | [Open](27_vit_patch_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/27_vit_patch_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/27_vit_patch_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 13 | GPT-2 Block | 🔴 Hard | [Open](13_gpt2_block.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/13_gpt2_block.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/13_gpt2_block.ipynb\" target=\"_blank\">Colab</a> | [Open](13_gpt2_block_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/13_gpt2_block_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/13_gpt2_block_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 28 | Mixture of Experts | 🔴 Hard | [Open](28_moe.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/28_moe.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/28_moe.ipynb\" target=\"_blank\">Colab</a> | [Open](28_moe_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/28_moe_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/28_moe_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "\n",
+    "### ⚙️ Training & Optimization\n",
+    "\n",
+    "| # | Problem | Difficulty | Template | Solution |\n",
+    "|:---:|---------|:----------:|:--------:|:--------:|\n",
+    "| 29 | Adam Optimizer | 🟡 Medium | [Open](29_adam.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/29_adam.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/29_adam.ipynb\" target=\"_blank\">Colab</a> | [Open](29_adam_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/29_adam_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/29_adam_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 30 | Cosine LR Scheduler | 🟡 Medium | [Open](30_cosine_lr.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/30_cosine_lr.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/30_cosine_lr.ipynb\" target=\"_blank\">Colab</a> | [Open](30_cosine_lr_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/30_cosine_lr_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/30_cosine_lr_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 40 | Linear Regression | 🟡 Medium | [Open](40_linear_regression.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/40_linear_regression.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/40_linear_regression.ipynb\" target=\"_blank\">Colab</a> | [Open](40_linear_regression_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/40_linear_regression_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/40_linear_regression_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "\n",
+    "### 🎯 Inference & Decoding\n",
+    "\n",
+    "| # | Problem | Difficulty | Template | Solution |\n",
+    "|:---:|---------|:----------:|:--------:|:--------:|\n",
+    "| 32 | Top-k / Top-p Sampling | 🟡 Medium | [Open](32_topk_sampling.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/32_topk_sampling.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/32_topk_sampling.ipynb\" target=\"_blank\">Colab</a> | [Open](32_topk_sampling_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/32_topk_sampling_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/32_topk_sampling_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 33 | Beam Search | 🟡 Medium | [Open](33_beam_search.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/33_beam_search.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/33_beam_search.ipynb\" target=\"_blank\">Colab</a> | [Open](33_beam_search_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/33_beam_search_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/33_beam_search_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 34 | Speculative Decoding | 🔴 Hard | [Open](34_speculative_decoding.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/34_speculative_decoding.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/34_speculative_decoding.ipynb\" target=\"_blank\">Colab</a> | [Open](34_speculative_decoding_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/34_speculative_decoding_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/34_speculative_decoding_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "\n",
+    "### 🔬 Advanced\n",
+    "\n",
+    "| # | Problem | Difficulty | Template | Solution |\n",
+    "|:---:|---------|:----------:|:--------:|:--------:|\n",
+    "| 35 | BPE Tokenizer | 🔴 Hard | [Open](35_bpe.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/35_bpe.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/35_bpe.ipynb\" target=\"_blank\">Colab</a> | [Open](35_bpe_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/35_bpe_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/35_bpe_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 36 | INT8 Quantization | 🔴 Hard | [Open](36_int8_quantization.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/36_int8_quantization.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/36_int8_quantization.ipynb\" target=\"_blank\">Colab</a> | [Open](36_int8_quantization_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/36_int8_quantization_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/36_int8_quantization_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 37 | DPO Loss | 🔴 Hard | [Open](37_dpo_loss.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/37_dpo_loss.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/37_dpo_loss.ipynb\" target=\"_blank\">Colab</a> | [Open](37_dpo_loss_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/37_dpo_loss_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/37_dpo_loss_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 38 | GRPO Loss | 🔴 Hard | [Open](38_grpo_loss.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/38_grpo_loss.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/38_grpo_loss.ipynb\" target=\"_blank\">Colab</a> | [Open](38_grpo_loss_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/38_grpo_loss_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/38_grpo_loss_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "| 39 | PPO Loss | 🔴 Hard | [Open](39_ppo_loss.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/templates/39_ppo_loss.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/39_ppo_loss.ipynb\" target=\"_blank\">Colab</a> | [Open](39_ppo_loss_solution.ipynb) · <a href=\"https://github.com/duoan/TorchCode/blob/master/solutions/39_ppo_loss_solution.ipynb\" target=\"_blank\">GitHub</a> · <a href=\"https://colab.research.google.com/github/duoan/TorchCode/blob/master/solutions/39_ppo_loss_solution.ipynb\" target=\"_blank\">Colab</a> |\n",
+    "\n",
+    "## Useful Commands\n",
+    "\n",
+    "```python\n",
+    "from torch_judge import check, hint, status\n",
+    "\n",
+    "status()                   # Progress dashboard\n",
+    "check(\"relu\")              # Judge your implementation\n",
+    "hint(\"causal_attention\")   # Get a hint\n",
+    "```"
    ]
   }
  ],
@@ -65,4 +161,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

templates/01_relu.ipynb

@@ -30,6 +30,20 @@
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},

templates/02_softmax.ipynb

@@ -30,6 +30,20 @@
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},

templates/03_linear.ipynb

@@ -26,6 +26,20 @@
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},

templates/04_layernorm.ipynb

@@ -32,6 +32,20 @@
    ],
    "outputs": []
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "metadata": {},

templates/05_attention.ipynb

@@ -29,6 +29,20 @@
     "- Must handle cross-attention (seq_q ≠ seq_k)"
    ]
   },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "code",
    "execution_count": null,

templates/06_multihead_attention.ipynb

@@ -37,6 +37,21 @@
     "5. Output projection: `self.W_o(concat)`"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "02a059c4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

templates/07_batchnorm.ipynb

@@ -1,131 +1,145 @@
 {
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "id": "89fd15cb",
-      "metadata": {},
-      "source": [
-        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/07_batchnorm.ipynb)\n",
-        "\n",
-        "# 🟡 Medium: Implement BatchNorm\n",
-        "\n",
-        "Implement **Batch Normalization** with both **training** and **inference** behavior.\n",
-        "\n",
-        "In training mode, use **batch statistics** and update running estimates:\n",
-        "\n",
-        "$$\\text{BN}(x) = \\gamma \\cdot \\frac{x - \\mu_B}{\\sqrt{\\sigma_B^2 + \\epsilon}} + \\beta$$\n",
-        "\n",
-        "where $\\mu_B$ and $\\sigma_B^2$ are the mean and variance computed **across the batch** (dim=0).\n",
-        "\n",
-        "In inference mode, use the provided **running mean/var** instead of current batch stats.\n",
-        "\n",
-        "### Signature\n",
-        "```python\n",
-        "def my_batch_norm(\n",
-        "    x: torch.Tensor,\n",
-        "    gamma: torch.Tensor,\n",
-        "    beta: torch.Tensor,\n",
-        "    running_mean: torch.Tensor,\n",
-        "    running_var: torch.Tensor,\n",
-        "    eps: float = 1e-5,\n",
-        "    momentum: float = 0.1,\n",
-        "    training: bool = True,\n",
-        ") -> torch.Tensor:\n",
-        "    # x: (N, D) — normalize each feature across all samples in the batch\n",
-        "    # running_mean, running_var: updated in-place during training; used as-is during inference\n",
-        "```\n",
-        "\n",
-        "### Rules\n",
-        "- Do **NOT** use `F.batch_norm`, `nn.BatchNorm1d`, etc.\n",
-        "- Compute batch mean and variance over `dim=0` with `unbiased=False`\n",
-        "- Update running stats like PyTorch: `running = (1 - momentum) * running + momentum * batch_stat`\n",
-        "- Use `running_mean` / `running_var` for inference when `training=False`\n",
-        "- Must support autograd w.r.t. `x`, `gamma`, `beta`（running statistics 应视作 buffer，而不是需要梯度的参数）"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "import torch"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "d946ca79",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# ✏️ YOUR IMPLEMENTATION HERE\n",
-        "\n",
-        "def my_batch_norm(\n",
-        "    x,\n",
-        "    gamma,\n",
-        "    beta,\n",
-        "    running_mean,\n",
-        "    running_var,\n",
-        "    eps=1e-5,\n",
-        "    momentum=0.1,\n",
-        "    training=True,\n",
-        "):\n",
-        "    pass  # Replace this"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "id": "26b93e71",
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# 🧪 Debug\n",
-        "x = torch.randn(8, 4)\n",
-        "gamma = torch.ones(4)\n",
-        "beta = torch.zeros(4)\n",
-        "\n",
-        "# Running stats typically live on the same device and shape as features\n",
-        "running_mean = torch.zeros(4)\n",
-        "running_var = torch.ones(4)\n",
-        "\n",
-        "# Training mode: uses batch stats and updates running_mean / running_var\n",
-        "out_train = my_batch_norm(x, gamma, beta, running_mean, running_var, training=True)\n",
-        "print(\"[Train] Output shape:\", out_train.shape)\n",
-        "print(\"[Train] Column means:\", out_train.mean(dim=0))   # should be ~0\n",
-        "print(\"[Train] Column stds: \", out_train.std(dim=0))    # should be ~1\n",
-        "print(\"Updated running_mean:\", running_mean)\n",
-        "print(\"Updated running_var:\", running_var)\n",
-        "\n",
-        "# Inference mode: uses running_mean / running_var only\n",
-        "out_eval = my_batch_norm(x, gamma, beta, running_mean, running_var, training=False)\n",
-        "print(\"[Eval] Output shape:\", out_eval.shape)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# ✅ SUBMIT\n",
-        "from torch_judge import check\n",
-        "check(\"batchnorm\")"
-      ]
-    }
-  ],
-  "metadata": {
-    "kernelspec": {
-      "display_name": "Python 3",
-      "language": "python",
-      "name": "python3"
-    },
-    "language_info": {
-      "name": "python",
-      "version": "3.11.0"
-    }
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "89fd15cb",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/duoan/TorchCode/blob/master/templates/07_batchnorm.ipynb)\n",
+    "\n",
+    "# 🟡 Medium: Implement BatchNorm\n",
+    "\n",
+    "Implement **Batch Normalization** with both **training** and **inference** behavior.\n",
+    "\n",
+    "In training mode, use **batch statistics** and update running estimates:\n",
+    "\n",
+    "$$\\text{BN}(x) = \\gamma \\cdot \\frac{x - \\mu_B}{\\sqrt{\\sigma_B^2 + \\epsilon}} + \\beta$$\n",
+    "\n",
+    "where $\\mu_B$ and $\\sigma_B^2$ are the mean and variance computed **across the batch** (dim=0).\n",
+    "\n",
+    "In inference mode, use the provided **running mean/var** instead of current batch stats.\n",
+    "\n",
+    "### Signature\n",
+    "```python\n",
+    "def my_batch_norm(\n",
+    "    x: torch.Tensor,\n",
+    "    gamma: torch.Tensor,\n",
+    "    beta: torch.Tensor,\n",
+    "    running_mean: torch.Tensor,\n",
+    "    running_var: torch.Tensor,\n",
+    "    eps: float = 1e-5,\n",
+    "    momentum: float = 0.1,\n",
+    "    training: bool = True,\n",
+    ") -> torch.Tensor:\n",
+    "    # x: (N, D) — normalize each feature across all samples in the batch\n",
+    "    # running_mean, running_var: updated in-place during training; used as-is during inference\n",
+    "```\n",
+    "\n",
+    "### Rules\n",
+    "- Do **NOT** use `F.batch_norm`, `nn.BatchNorm1d`, etc.\n",
+    "- Compute batch mean and variance over `dim=0` with `unbiased=False`\n",
+    "- Update running stats like PyTorch: `running = (1 - momentum) * running + momentum * batch_stat`\n",
+    "- Use `running_mean` / `running_var` for inference when `training=False`\n",
+    "- Must support autograd w.r.t. `x`, `gamma`, `beta`（running statistics 应视作 buffer，而不是需要梯度的参数）"
+   ]
   },
-  "nbformat": 4,
-  "nbformat_minor": 5
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Install torch-judge in Colab (no-op in JupyterLab/Docker)\n",
+    "try:\n",
+    "    import google.colab\n",
+    "    get_ipython().run_line_magic('pip', 'install -q torch-judge')\n",
+    "except ImportError:\n",
+    "    pass\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d946ca79",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ✏️ YOUR IMPLEMENTATION HERE\n",
+    "\n",
+    "def my_batch_norm(\n",
+    "    x,\n",
+    "    gamma,\n",
+    "    beta,\n",
+    "    running_mean,\n",
+    "    running_var,\n",
+    "    eps=1e-5,\n",
+    "    momentum=0.1,\n",
+    "    training=True,\n",
+    "):\n",
+    "    pass  # Replace this"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "26b93e71",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 🧪 Debug\n",
+    "x = torch.randn(8, 4)\n",
+    "gamma = torch.ones(4)\n",
+    "beta = torch.zeros(4)\n",
+    "\n",
+    "# Running stats typically live on the same device and shape as features\n",
+    "running_mean = torch.zeros(4)\n",
+    "running_var = torch.ones(4)\n",
+    "\n",
+    "# Training mode: uses batch stats and updates running_mean / running_var\n",
+    "out_train = my_batch_norm(x, gamma, beta, running_mean, running_var, training=True)\n",
+    "print(\"[Train] Output shape:\", out_train.shape)\n",
+    "print(\"[Train] Column means:\", out_train.mean(dim=0))   # should be ~0\n",
+    "print(\"[Train] Column stds: \", out_train.std(dim=0))    # should be ~1\n",
+    "print(\"Updated running_mean:\", running_mean)\n",
+    "print(\"Updated running_var:\", running_var)\n",
+    "\n",
+    "# Inference mode: uses running_mean / running_var only\n",
+    "out_eval = my_batch_norm(x, gamma, beta, running_mean, running_var, training=False)\n",
+    "print(\"[Eval] Output shape:\", out_eval.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ✅ SUBMIT\n",
+    "from torch_judge import check\n",
+    "check(\"batchnorm\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
 }