{
 "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",
   "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/11_sliding_window.ipynb)\n",
    "\n",
    "# 🔴 Hard: Sliding Window Attention\n",
    "\n",
    "Implement **Sliding Window Attention** — used in Longformer, Mistral, etc. for efficient long-context processing.\n",
    "\n",
    "Each position $i$ can only attend to positions $j$ where $|i - j| \\le w$ (the window size).\n",
    "\n",
    "### Signature\n",
    "```python\n",
    "def sliding_window_attention(Q, K, V, window_size):\n",
    "    # Q, K, V: (batch, seq, d) → output: (batch, seq, d_v)\n",
    "    # window_size: int — position i attends to [i-w, i+w]\n",
    "```\n",
    "\n",
    "### Rules\n",
    "- Do **NOT** use sparse attention libraries\n",
    "- Mask positions outside the window with `-inf`\n",
    "- `window_size=0`: only self — output should equal V\n",
    "- `window_size >= seq_len`: equivalent to full attention"
   ],
   "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": {},
   "source": [
    "import torch\n",
    "import math"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# ✏️ YOUR IMPLEMENTATION HERE\n",
    "\n",
    "def sliding_window_attention(Q, K, V, window_size):\n",
    "    pass  # Replace this"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# 🧪 Debug\n",
    "Q = torch.randn(1, 6, 8)\n",
    "K = torch.randn(1, 6, 8)\n",
    "V = torch.randn(1, 6, 8)\n",
    "\n",
    "out = sliding_window_attention(Q, K, V, window_size=1)\n",
    "print(\"Output shape:\", out.shape)  # (1, 6, 8)\n",
    "\n",
    "# window=0 should return V\n",
    "out0 = sliding_window_attention(Q, K, V, window_size=0)\n",
    "print(\"window=0 == V?\", torch.allclose(out0, V, atol=1e-5))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "from torch_judge import check\n",
    "check('sliding_window')"
   ],
   "outputs": [],
   "execution_count": null
  }
 ]
}