Created using Colab

chapter1_transformer_interp/exercises/part1_transformer_from_scratch/1.1_Transformer_from_Scratch_exercises.ipynb

@@ -2592,16 +2592,46 @@
         "> You should spend up to 10-15 minutes on this exercise.\n",
         "> ```\n",
         "\n",
-        "Now, we can put together the attention, MLP and layernorms into a single transformer block. Remember to implement the residual connections correctly!"
+        "Now, we can put together the attention, MLP and layernorms into a single transformer block. Remember to implement the residual connections correctly!\n",
+        "\n",
+        "![image.png](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAZUAAAL1CAIAAACAJe6rAAAAAXNSR0IArs4c6QAAAERlWElmTU0AKgAAAAgAAYdpAAQAAAABAAAAGgAAAAAAA6ABAAMAAAABAAEAAKACAAQAAAABAAABlaADAAQAAAABAAAC9QAAAAA42QnTAABAAElEQVR4AexdCVhV1dreaooKCoKIOIKCA+CAoqQ5oTmLOOTQVdMyzbw5dNWC9E/tqlBppna9SjmVlsN1AJzSHEpzNpFATVFwRCRwAsWh/N/DgsVmn4HDOfvA2Wd/++HZrL3Wt7611rv2fs+35lIvXrwQ6CIECAFCQIEIlFZgninLhAAhQAhoECD+oveAECAElIoA8ZdSa47yTQgQAsRf9A4QAoSAUhEg/lJqzVG+CQFCgPiL3gFCgBBQKgLEX0qtOco3IUAIEH/RO0AIEAJKRYD4S6k1R/kmBAgB4i96BwgBQkCpCBB/KbXmKN+EACFA/EXvACFACCgVAeIvpdYc5ZsQIASIv+gdIAQIAaUiQPyl1JqjfBMChADxF70DhAAhoFQEiL+UWnOUb0KAECD+oneAECAElIoA8ZdSa47yTQgQAsRf9A4QAoSAUhEg/lJqzVG+CQFCgPiL3gFCgBBQKgLEX0qtOco3IUAIEH/RO0AIEAJKRYD4S6k1R/kmBAgB4i96BwgBQkCpCBB/KbXmKN+EACFA/EXvACFACCgVAeIvpdYc5ZsQIASIv+gdIAQIAaUiQPyl1JqjfBMChADxF70DhAAhoFQEiL+UWnOUb0KAECD+oneAECAElIoA8ZdSa47yTQgQAsRf9A4QAoSAUhF4SakZ15PvZ8+e6Qkhb0KAECiAQNmyZQs8K/BBSfz16NGjjIwMgPz48ePdu3dfuXL5xImTR44cUSDslGVCwOoQ8PP1DXz5ZS8v74EDB1SoUAH5c3BwcHJysrqMijJU6sWLF6JHa3TeuHHj0qVLaWlpR48e+WnvT/EJCeJcAnTPeh7VqlWvVMnB3t5BHERuQoAQMIxAVlbmw4eZd+7cTk+/m5SUlJKSwuVdKju07dhx+PA3XF1da9Wq5e3tzYOsx2G9/AXa2rVr17lzCWIjC5j2GzQkMDDQ19cXIAJZ3PFbgR8KOMqVK2c9yFJOCAHrR+Dp06fIZGZmJto0cMBKwD0hIQHfndhWYKZZs2bNevToYVVEZo38FRMTExERIW4Yjh49OiQkpEuXLmixo4cLd1wAWtzbJXYjiC5CgBAwBgH2KTFJ/lmxrywrK+vcuXObNm38buWK9AeZTMbT0zMsLGzMmDHGKLe0jLXwF/BKTk5Gr9bEiRNZmYODezdt2rxXr15t27ZlPpDBZWlESD8hQAgwBDRmQl4fP/pwNm/ecuTIoVOnfmPNzMmTJ40f/8+aNWtWrFixpBCzCv4CNGvWrJk7dy5QQAtx5px5fn5+/v7+rO8Q3fYlhQ6lSwgQAgwBxmUwIE6ePIkG5urVq1kLadiwYWgetWvXjjNdcSJWwvwF5vr0009XrFiBMqONPS88vFGjRqyBTbRVnO8BpUUIGIkAs7bYqNqMGTMYi6G19P77U4KCgoxUIpdYifEX6AlNRcZcKMz+/ftReNY8pEaiXLVLeggBCyHAzDEoB39NmzaNs9i3364tzikXJcBf9+7d27Nnz5AhQ1B40HZwcMioUaMABxlcFnrVSC0hYDkEmDl24MCBdevWMXMkPDxi+PBhmHJhuUS55uLmLzQYZ8+ejaIiB5GRkT179kQ5ibl4fZCDEFAiAmAxNJsw/sZmDqAvaPGSJcXQnCwza9asYsPr66+/7tat2++//44+v507d6J4mLr15MmTYssAJUQIEAKWQADkVaZMmcaNG3fv3h1ctnnzZozIPX/+vH379vC3RIpMZzHZX+jtGz9+XEzMDqQaHR0dHByMAuOyXMFIMyFACBQ/AqxfDN1hr7zyCkv94sWLlpvyWhz7T6DNWLt2bZAXxlmvX78O8kKDkcir+N8tSpEQsDQC+K7xdWPOJua+Tp8+Hck1aNAAvWMW+t4tbn9hMn3fvn1RjMWLF48YMQJjE9TbZel3iPQTAiWOAJsOhh4x9vmjU3/KlH8xTxnzZln+WrJkCZtPj+kRWLSI3FuIhmVEhFQRAoSALAgwtkpNTUXzCwrR67127VpZNHMlFuSvt99+G+OpWC31yy+/0CAjR5wchICqEEB3PppcQ4cOZt3fWDEuoxVmkf4vGFnvvz8Z5IVmMJGXql5WKiwhIEEA5AXCWr9+I9ZLImjgwP4YzZPImPwo//6FIK8PPpj25ZeLQF4bNmwgy8vkuqGIhIBtIABOAIXNnDkLxQEzCMI40BnsMvNLJ3/7EZZXThYFDEAw09H8XJIGQoAQUDoCoDAQGVs16O7ufuvWLfNLJHP7Ees5meUF8kJ2zRxqhAZ+mV9U0kAIEAIliACzwr755hs0JLEDTxM/P/MzI6f9tXHjRqxqNL7ZCG566dG1DT9EJd/TPZHV3r5i+fLlHR0dq7Xs+kodzQ6rgMD8Mpum4aflYeee1MjMfPjunDlVCtvSB4Zn3MZ/77yi2Q/W279Pl8Ca+vaGBQh/XtyxJuYPSDo0G/pez7pmkr5ppZMxFkrEtnKbMGGC0ssiIyykiiGA14NbYZggNmfOHHOQkY2/MEWtc+fOyMqvv/4KCjPmxWWfbge//omFlaBNm55HffseCB0CFisRCrOr+OIjF4fPNIeHCP979qLvC+TCEJNCPmHjF82GfJxTMr8lO7frIyaAcOPsv+u1+jckfaafTpjTwhjoctRa6Q3cDTMcu7nBBrfSLFK2ShQBRmH29vbIBbrIBw8ebHJ25Gk/YkCBkRfmeRlJXjzHNfJcmGmBVrHkwnaGCD96dJfwzbtBYyJ+vZaJwufFKNb/lZyLnJxmc37NFT/hzfkHL98tqZyzTBTPHWXEcouobduQ3OLFS2Tpoy2enFMqxYYAfvvxYmApDlJEiw0vjMlJy8NfrVu3Rg6wnwSWZJtoPviEHrpyBV16kivtWmzUFxO9WPkOfLbkuz0mF7UkI6Z+FTRD80mr4cIuw+yMqHVrv8OrqQbWVkO1yltGsARmJqCtBrVYYGQiaQiCDPyFSfbojUOfHLbxMtyqKhQCFENyPa9Yp+/7izZv+ITF3bL5xxI0wQrNvyGB9W9N3n3Xtu0RUBV6vsLCQhkOYDEsHzGECYWpGAF86c2bN8eyQmAwduxY05Awl7/wA4sBUWz3M2XKVOTATP7SLgMU4moS3HqAZgUCmmJPz+cc8QQnvhbtCwQBzxzR3BsemScT5m6xjCXdbv26/6NFTgJLPwq17VYkEIbxJQYzcvlyMJqkRsQC5FYzAngxxo0bhx4nbAiI0T8ToDCLv8CgU6a8j1SxVxmsQRCNCTkoNArU/l2hXSP3yjmSZ3FHsSs8S1kzyRuDeuVeWVm67NOF7/XQuMuVK1XKc/lPt9gHw+7nt0YMbVWXheJer17zoXO3FKcR12fQOzPDvDWZj18VFLHhTkF6zSmULdyANt4HbnyxIjETjFWELRSSyiArAowxduzQbKtlWkeYWfyFFUJY04SWo+ndXsbBUUa4kCdYPs+R9//x8Y9fd/gwcj96+nM6+5MPHD+GMHwzmJzx+VB3DAJuOHWNheKelHR2y+yhQV4dOM3lKbLQ/9Tfsyu2GvZZrv34zaytu67a5PeMQk2f/pE2iDDPwWs2WWTtwpJPUREAhWFPGnSdIyJ2Zi5qdNP5i7UckR5ajhayvFhhNH1GZ37ffeIBHj069Gns6ipOzuXapnnrBa8B01et/R5/mFEyYUQ3Rl6fjGgxK0YTy3nge0u/XrFpW7RGZvbonNGA+Am9+oDCNMotf7n79ps5/5OcdFIn9HqnOK0/yxdOkwIAx/uAU051JoffOeIvnciQJxDAzxu6zrE7BVqR2G6rSJiYvv4RSSIlbKYq4wpH7bccPi/uX/5k5qTfcorVIqiRpHg4Fhjk9fVnUzrVr4IgnNHNBM59O4WRl/D2f68vegN6cGmCgoNbNKo7e+rHW67HT5g11a9BZPHMKfMd/K9/T//4/zRT3X7EWOSLH94UszDLs3LvwBY9X/yIZklBYIJh60rI2FKRJWWkR3MQwLsxfvx48Bc2C3vx4oXxqky0v/Bjiy1iQZnY31qel7J06v9+2ny44LV169bP3ulS2skrl4mEgH+08QNPSVLs3bkfPMHiuPJKfrrfxCiN2yf0xtfjGHMhFruaDv6/D0d214SeOBV/UWOgFcP197NyMy6d6s9SWv/WP2PuFI/pVwxFA7x4HyQ9X5J0MReM1YLEnx4JASCALxe9+NjjEG7MZzAeExPtrwH9NV8iflRl2081ftXkrqv05dvRUbh/P+B/u1cP7O4rIikm7vGsTh1xRA0v/LiezenvNC7kxY0bqeLgHHfrAUOFOT8Kgqaz7K2OwfDT93WB8rRim+KRo6fprA2fHB7ycZogLO37xqDEH9iiKFPUWVkcybCjdu4wF2zgwAEeHh5y4amdBPkoGgG8GDh1Db+CCxcu7N+/P1p1xhTHFP5CGxXjSji6EYeGa7GJMYnqlsHMe+2Av7Ket+jQq1e3ljUadNZFXohR1c/13rNnVcRxEy7mzuiN2xoxfu9zcRBzezlkM8eWPx5NSHnS2t0Ow/zaYjgeycHBIWehw1PtUBN8mnT/x9SxBzHaoGlFRmzIWvSG0q0w8D7WCYmNL0ymyXr0KCkpCb+of8THsUYlG4jEYLkJoFEUNSAA/gJnof8aK8/Q8MLiWWNKXWT+QjJYsgTVoaEf4d2V7efUJ/T76A+87HSsmAODYN4DvnM9XJmNHn1JUbPu5NJTxoEoff2BGIvk/TX3k/a87df/kru7i7MzV5Vy/Wp6y/FLpo1/r6fTX4IM/AWsyjrW/2DciOO79m/B2olvVqwZ3fndlxvwFJXowDswf/58lnNAOuKt0TNLine truncated
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 200,
       "metadata": {
-        "id": "4Qwi3GknfdfP"
+        "colab": {
+          "base_uri": "https://localhost:8080/"
         },
-      "outputs": [],
+        "id": "4Qwi3GknfdfP",
+        "outputId": "fdc68f75-a896-4ca4-fc84-539cd8e3179d"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Input shape: torch.Size([2, 4, 768])\n",
+            "torch.Size([2, 4, 768])\n",
+            "torch.Size([12, 768, 64])\n",
+            "torch.Size([12, 768, 64])\n",
+            "torch.Size([12, 768, 64])\n",
+            "Output shape: torch.Size([2, 4, 768]) \n",
+            "\n",
+            "Input shape: torch.Size([1, 35, 768])\n",
+            "torch.Size([1, 35, 768])\n",
+            "torch.Size([12, 768, 64])\n",
+            "torch.Size([12, 768, 64])\n",
+            "torch.Size([12, 768, 64])\n",
+            "Output shape: torch.Size([1, 35, 768])\n",
+            "Reference output shape: torch.Size([1, 35, 768]) \n",
+            "\n",
+            "100.00% of the values are correct\n",
+            "\n"
+          ]
+        }
+      ],
       "source": [
         "class TransformerBlock(nn.Module):\n",
         "    def __init__(self, cfg: Config):\n",
@@ -2613,7 +2643,10 @@
         "        self.mlp = MLP(cfg)\n",
         "\n",
         "    def forward(self, resid_pre: Float[Tensor, \"batch position d_model\"]) -> Float[Tensor, \"batch position d_model\"]:\n",
-        "        raise NotImplementedError()\n",
+        "        l1 = self.ln1.forward(resid_pre)\n",
+        "        l2 = self.attn.forward(l1) + resid_pre\n",
+        "        l3 = self.ln2.forward(l2)\n",
+        "        return self.mlp.forward(l3)+l2\n",
         "\n",
         "\n",
         "tests.rand_float_test(TransformerBlock, [2, 4, 768])\n",