cut_cross_entropy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0ef58be0-882f-4e0e-afac-8b69ef58fd56",
   "metadata": {},
   "source": [
    "# cut cross entropy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1130a1e1-1aeb-473f-804a-476baaefc344",
   "metadata": {},
   "source": [
    "author: [xiaodongguaAIGC](https://github.com/dhcode-cpp)\n",
    "reference: [CCE](https://arxiv.org/abs/2411.09009)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "78c71de7-0c5f-4637-a06f-d44c47ccce25",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x10af36bb0>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "torch.manual_seed(42)   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7db6c00-65a1-4c9c-987c-475043be7d3e",
   "metadata": {},
   "source": [
    "## Cross-Entropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c9cef966-b55c-4dda-8122-504b6c87a0f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.tensor([1.0, 3.0, -1.2, 1.1, -0.5, -0.8]) # logits\n",
    "Y = torch.tensor([0.4, 0.2, 0.1,  0.1, 0.05, 0.15])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cde6066-bd74-493f-ab92-f3f8ad4d53bf",
   "metadata": {},
   "source": [
    "### Cross Entropy Pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "338ea46f-84dc-4f63-9e27-d91d04b99598",
   "metadata": {},
   "source": [
    "$CE(X,Y) = \\sum_i -\\log(\\text{softmax}(X))$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a1567690-e4a2-42b6-8bc4-4385a00ccbad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.4569</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2.4569\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "loss = nn.CrossEntropyLoss()\n",
    "l = loss(X,Y)\n",
    "print(l)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eef0206d-a586-4883-84f9-9357a4d68d66",
   "metadata": {},
   "source": [
    "### Cross Entropy with log softmax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b8c834d-3318-4888-bfbf-4f7321ae1103",
   "metadata": {},
   "source": [
    "when we calculate cross entropy, we need logprob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6a294152-b93c-45d9-a5b4-dc147c56daa1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.4569</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2.4569\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "l = -Y * F.log_softmax(X, dim = 0 )\n",
    "l = l.sum()\n",
    "print(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a5f45b72-508f-4c87-916c-061fac126af9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.4569</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2.4569\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "p = F.softmax(X, dim = 0 )\n",
    "l = -Y * p.log()\n",
    "l = l.sum()\n",
    "print(l)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "211365fd-4f30-4361-8f10-aa31eb88f41f",
   "metadata": {},
   "source": [
    "## Softmax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59a85d99-39ea-438f-902b-fdf08cf3e7a8",
   "metadata": {},
   "source": [
    "### forward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2f20c2c0-3597-403a-b4b1-9a46970f8418",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1001</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7394</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0111</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1106</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0223</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0165</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0.1001\u001b[0m, \u001b[1;36m0.7394\u001b[0m, \u001b[1;36m0.0111\u001b[0m, \u001b[1;36m0.1106\u001b[0m, \u001b[1;36m0.0223\u001b[0m, \u001b[1;36m0.0165\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.4569</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2.4569\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "p = torch.exp(X) / (torch.exp(X)).sum()\n",
    "print(p)\n",
    "l = - Y * p.log()\n",
    "l = l.sum()\n",
    "print(l)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c7b46e8-ea31-4c7e-8ed6-2ef338ea7c3d",
   "metadata": {},
   "source": [
    "### gradient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c5c782ee-e9ee-4952-bad2-944bcf68c64a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0500</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0899</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0148</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0018</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0046</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0222</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-0.0500\u001b[0m,  \u001b[1;36m0.0899\u001b[0m, \u001b[1;36m-0.0148\u001b[0m,  \u001b[1;36m0.0018\u001b[0m, \u001b[1;36m-0.0046\u001b[0m, \u001b[1;36m-0.0222\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "N = 6 \n",
    "dX = (p - Y) / 6\n",
    "print(dX)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "36d64972-236c-4699-8c4a-3e9f57022eee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.2999</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.5394</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0889</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0106</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0277</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1335</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-0.2999\u001b[0m,  \u001b[1;36m0.5394\u001b[0m, \u001b[1;36m-0.0889\u001b[0m,  \u001b[1;36m0.0106\u001b[0m, \u001b[1;36m-0.0277\u001b[0m, \u001b[1;36m-0.1335\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dpdX = torch.diag(p) - torch.outer(p, p)\n",
    "dX = -(Y / p) @ dpdX   # first term is (- Y log p)' = (-Y/p)\n",
    "print(dX)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d6d2f7b-ff56-462a-95b4-1faef28b6f39",
   "metadata": {},
   "source": [
    "## Log Softmax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45685c53-2f6c-448b-af12-716ecc012809",
   "metadata": {},
   "source": [
    "### forward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e08e8380-3559-498d-ba8b-38a19859809a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.3019</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.3019</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.5019</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.2019</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.8019</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.1019</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-2.3019\u001b[0m, \u001b[1;36m-0.3019\u001b[0m, \u001b[1;36m-4.5019\u001b[0m, \u001b[1;36m-2.2019\u001b[0m, \u001b[1;36m-3.8019\u001b[0m, \u001b[1;36m-4.1019\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.4569</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2.4569\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "LSE = torch.log( torch.sum( torch.exp( X ) ) )\n",
    "p = X - LSE # X - LSE\n",
    "print(p)\n",
    "l = - Y * p # delete log\n",
    "l = l.sum()\n",
    "print(l)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fce234df-f75a-4cef-a231-156e96b59a97",
   "metadata": {},
   "source": [
    "### gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23914820-a8b0-4e2c-b913-c69c60002f2c",
   "metadata": {},
   "source": [
    "首先，L = -∑y_i * log(softmax(x_i))\n",
    "\n",
    "= -∑y_i * p_i (因为p_i = LogSoftmax(x_i))\n",
    "\n",
    "求导链式法则：\n",
    "\n",
    "∂L/∂x_k = ∑_i (∂L/∂p_i * ∂p_i/∂x_k)\n",
    "\n",
    "其中：\n",
    "\n",
    "∂L/∂p_i = -y_i\n",
    "\n",
    "∂p_i/∂x_k = δ_{ik} - softmax(x_k) \n",
    "\n",
    "代入得到：\n",
    "\n",
    "∂L/∂x_k = ∑_i (-y_i * (δ_{ik} - softmax(x_k)))\n",
    "\n",
    "= -y_k + ∑_i y_i * softmax(x_k)\n",
    "\n",
    "= -y_k + softmax(x_k) * ∑_i y_i\n",
    "\n",
    "= softmax(x_k) - y_k (因为∑y_i = 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a09fb378-d88f-418f-b907-142039bf95bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.2999</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.5394</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0889</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0106</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0277</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1335</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-0.2999\u001b[0m,  \u001b[1;36m0.5394\u001b[0m, \u001b[1;36m-0.0889\u001b[0m,  \u001b[1;36m0.0106\u001b[0m, \u001b[1;36m-0.0277\u001b[0m, \u001b[1;36m-0.1335\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "l = - Y * p / N\n",
    "# p.exp() is softmax result probability\n",
    "dX = p.exp() - Y # lse is easy revearse p \n",
    "print(dX)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64d2a8ef-0602-448d-bac3-e15fe410c327",
   "metadata": {},
   "source": [
    "## Cut Cross Entropy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a047fcc-e378-47e4-bb2a-67c5434c431c",
   "metadata": {},
   "source": [
    "think about llm model has a big vocabulary, so the classifer project d-dimension to |V|-dimension, and |V| usually is big like 128,000 or 256,000.\n",
    "\n",
    "classifier proj is big GEMM when the model params < 2B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "463a8255-adcf-40c1-9144-a9c3e2caaa08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;36m2\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;36m4\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;36m32\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "V = 128\n",
    "D = 16\n",
    "BS = 8\n",
    "block_size = 4\n",
    "n_b = BS // block_size\n",
    "d_b = D // block_size\n",
    "v_b = V // block_size\n",
    "\n",
    "print(n_b)\n",
    "print(d_b)\n",
    "print(v_b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ab8ee0f6-11ab-4a42-be5a-548f3d0f410b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m128\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "C = torch.randn(D, V)\n",
    "E = torch.randn(BS, D)\n",
    "Y_raw = torch.randn(BS, V)\n",
    "Y_t = Y_raw / 0.0001 # be sharp\n",
    "Y = F.softmax(Y_t, dim = 1)\n",
    "# # print(YY)\n",
    "# _, label = torch.max(Y, dim = 1)\n",
    "# print(label)\n",
    "print(Y.shape) # label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4f5cd23-a75d-4af0-b9f0-ee77928c3640",
   "metadata": {},
   "source": [
    "### standard CE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3c927951-75c8-4ec9-8e50-93a70e6ad299",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m128\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m128\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">12.3959</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m12.3959\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# ce\n",
    "logits = E @ C \n",
    "print(logits.shape)\n",
    "print(Y.shape)\n",
    "l = loss(logits, Y)\n",
    "print(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0ce001b1-21ef-4d14-b32f-47dddfd93688",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.1829e-02</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.1995e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4.3568e-04</span>,  <span style=\"color: #808000; text-decoration-color: #808000\">...</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.7271e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5.7252e-06</span>,\n",
       "         <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4.6990e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4.9314e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.7518e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.8711e-07</span>,  <span style=\"color: #808000; text-decoration-color: #808000\">...</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.5461e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.6502e-03</span>,\n",
       "         <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5.9683e-06</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.3456e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.4222e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4.7783e-02</span>,  <span style=\"color: #808000; text-decoration-color: #808000\">...</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.5149e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4.6281e-05</span>,\n",
       "         <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.3086e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"color: #808000; text-decoration-color: #808000\">...</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.8751e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0904e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7841e-04</span>,  <span style=\"color: #808000; text-decoration-color: #808000\">...</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.5724e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.1881e-04</span>,\n",
       "         <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5.2249e-04</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7716e-10</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.9223e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.3995e-07</span>,  <span style=\"color: #808000; text-decoration-color: #808000\">...</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2832e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.3845e-07</span>,\n",
       "         <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.8057e-07</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.2445e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.8148e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.2590e-05</span>,  <span style=\"color: #808000; text-decoration-color: #808000\">...</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.5709e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.8583e-05</span>,\n",
       "         <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.3119e-04</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m2.1829e-02\u001b[0m, \u001b[1;36m6.1995e-06\u001b[0m, \u001b[1;36m4.3568e-04\u001b[0m,  \u001b[33m...\u001b[0m, \u001b[1;36m6.7271e-05\u001b[0m, \u001b[1;36m5.7252e-06\u001b[0m,\n",
       "         \u001b[1;36m4.6990e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m4.9314e-03\u001b[0m, \u001b[1;36m6.7518e-06\u001b[0m, \u001b[1;36m2.8711e-07\u001b[0m,  \u001b[33m...\u001b[0m, \u001b[1;36m3.5461e-07\u001b[0m, \u001b[1;36m1.6502e-03\u001b[0m,\n",
       "         \u001b[1;36m5.9683e-06\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m2.3456e-05\u001b[0m, \u001b[1;36m3.4222e-04\u001b[0m, \u001b[1;36m4.7783e-02\u001b[0m,  \u001b[33m...\u001b[0m, \u001b[1;36m1.5149e-07\u001b[0m, \u001b[1;36m4.6281e-05\u001b[0m,\n",
       "         \u001b[1;36m3.3086e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[33m...\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m1.8751e-03\u001b[0m, \u001b[1;36m1.0904e-07\u001b[0m, \u001b[1;36m1.7841e-04\u001b[0m,  \u001b[33m...\u001b[0m, \u001b[1;36m1.5724e-05\u001b[0m, \u001b[1;36m1.1881e-04\u001b[0m,\n",
       "         \u001b[1;36m5.2249e-04\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m1.7716e-10\u001b[0m, \u001b[1;36m2.9223e-05\u001b[0m, \u001b[1;36m3.3995e-07\u001b[0m,  \u001b[33m...\u001b[0m, \u001b[1;36m3.2832e-06\u001b[0m, \u001b[1;36m9.3845e-07\u001b[0m,\n",
       "         \u001b[1;36m1.8057e-07\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m2.2445e-04\u001b[0m, \u001b[1;36m9.8148e-06\u001b[0m, \u001b[1;36m1.2590e-05\u001b[0m,  \u001b[33m...\u001b[0m, \u001b[1;36m6.5709e-05\u001b[0m, \u001b[1;36m3.8583e-05\u001b[0m,\n",
       "         \u001b[1;36m6.3119e-04\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0968</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0.0968\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# softmax -> ce\n",
    "P = F.softmax(logits, dim = 1)\n",
    "print(P)\n",
    "l = -Y * P.log()\n",
    "print(l.sum() / (BS*V))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b0219dbb-7073-4779-b199-b7a2f5af4f8e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0968</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0.0968\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# logsoftmax -> ce\n",
    "P = F.log_softmax(logits, dim = 1)\n",
    "l = -Y * P\n",
    "print(l.sum() / (BS*V))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "389a11ce-a84d-41bd-a347-d503c3981538",
   "metadata": {},
   "source": [
    "### Cut Cross Entropy "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dda1beb-4fa9-4020-a968-25645cf62ca3",
   "metadata": {},
   "source": [
    "#### *Algorithm 1: forward GEMM*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89d0676c-e398-4cbc-8946-30cf080fabfc",
   "metadata": {},
   "source": [
    "our target is get  \n",
    "\n",
    "\n",
    "logprob = X @ C - log \\sum_j X[i,:] C[:,j] # row sub\n",
    "\n",
    "the X[i,1] C[1,j] + X[i,2] C[2,j] + X[i,d] C[d,j] + \n",
    "\n",
    "output = -Y * logprob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "223d1885-bf07-46f8-a6ce-4ed06154ad92",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = E @ C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b36328fb-fe3d-4371-8e0c-7a7ca1c58f7f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m128\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(X.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5eb3cdd4-5873-440b-9c40-4f1d68d86282",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">64</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m64\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;36m2\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_blocks = torch.split(X, 64, dim = 1)\n",
    "print(X_blocks[0].shape)\n",
    "print(len(X_blocks))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "84754f0a-afc9-4f4f-9557-c954d68adeba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;36m4\u001b[0m \u001b[1;36m4\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m2\u001b[0m, \u001b[1;36m16\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m16\u001b[0m, \u001b[1;36m32\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "E_blocks = torch.split(E, n_b, dim = 0)\n",
    "C_blocks = torch.split(C, v_b, dim = 1)\n",
    "\n",
    "print(len(E_blocks), len(C_blocks))\n",
    "print(E_blocks[0].shape)\n",
    "print(C_blocks[0].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4d61446b-1ac8-4a06-9179-01035c8028f7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.2235</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.9430</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.3094</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.8064</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9664</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.5531</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.9595</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1982</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.5044</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.8647</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.5816</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0988</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.0377</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.3787</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.7961</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.4435</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.2511</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.4568</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4.3576</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.2228</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.9545</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.2834</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.8898</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.6405</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.7475</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m \u001b[1;36m6.2235\u001b[0m, \u001b[1;36m-1.9430\u001b[0m,  \u001b[1;36m2.3094\u001b[0m, \u001b[1;36m-0.8064\u001b[0m,  \u001b[1;36m0.9664\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.5531\u001b[0m,  \u001b[1;36m2.9595\u001b[0m, \u001b[1;36m-0.1982\u001b[0m, \u001b[1;36m-1.5044\u001b[0m, \u001b[1;36m-2.8647\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.5816\u001b[0m,  \u001b[1;36m1.0988\u001b[0m,  \u001b[1;36m6.0377\u001b[0m,  \u001b[1;36m1.3787\u001b[0m,  \u001b[1;36m2.7961\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.4435\u001b[0m, \u001b[1;36m-4.2511\u001b[0m, \u001b[1;36m-3.4568\u001b[0m,  \u001b[1;36m4.3576\u001b[0m, \u001b[1;36m-2.2228\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m2.9545\u001b[0m, \u001b[1;36m-5.2834\u001b[0m, \u001b[1;36m-3.8898\u001b[0m, \u001b[1;36m-3.6405\u001b[0m, \u001b[1;36m-6.7475\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "stand_mm = E @ C\n",
    "print(stand_mm[:5,:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "c1ca6872-626d-4d95-9d3d-d9d04326d010",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m128\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.2235</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.9430</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.3094</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.8064</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9664</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.5531</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.9595</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1982</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.5044</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.8647</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.5816</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0988</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.0377</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.3787</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.7961</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.4435</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.2511</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.4568</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4.3576</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.2228</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.9545</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.2834</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.8898</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.6405</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.7475</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m \u001b[1;36m6.2235\u001b[0m, \u001b[1;36m-1.9430\u001b[0m,  \u001b[1;36m2.3094\u001b[0m, \u001b[1;36m-0.8064\u001b[0m,  \u001b[1;36m0.9664\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.5531\u001b[0m,  \u001b[1;36m2.9595\u001b[0m, \u001b[1;36m-0.1982\u001b[0m, \u001b[1;36m-1.5044\u001b[0m, \u001b[1;36m-2.8647\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.5816\u001b[0m,  \u001b[1;36m1.0988\u001b[0m,  \u001b[1;36m6.0377\u001b[0m,  \u001b[1;36m1.3787\u001b[0m,  \u001b[1;36m2.7961\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.4435\u001b[0m, \u001b[1;36m-4.2511\u001b[0m, \u001b[1;36m-3.4568\u001b[0m,  \u001b[1;36m4.3576\u001b[0m, \u001b[1;36m-2.2228\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m2.9545\u001b[0m, \u001b[1;36m-5.2834\u001b[0m, \u001b[1;36m-3.8898\u001b[0m, \u001b[1;36m-3.6405\u001b[0m, \u001b[1;36m-6.7475\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CCE : Algorithm 1\n",
    "\n",
    "O = torch.zeros(BS, V)\n",
    "O_blocks = torch.split(O, n_b, dim = 0)\n",
    "O_result = []\n",
    "\n",
    "for e, o in zip(E_blocks, O_blocks): # i\n",
    "    e_d_blocks = torch.split(e, d_b, dim = 1)\n",
    "    c_blocks = torch.split(C, d_b, dim = 0)\n",
    "    for ed, cd in zip(e_d_blocks, c_blocks): # \n",
    "        a = ed @ cd\n",
    "        o = o + a\n",
    "    O_result.append(o)\n",
    "\n",
    "O_mat = torch.cat(O_result, dim = 0)\n",
    "print(O_mat.shape)\n",
    "\n",
    "print(O_mat[:5,:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2b0f833-3259-4af9-b046-797c67231c58",
   "metadata": {},
   "source": [
    "#### *Algorithm 2: forward logsoftmax*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca6a3395-cc9c-4212-907d-4fbe76cdf18f",
   "metadata": {},
   "source": [
    "Original LSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "251b84f1-c409-4750-8224-966cdc3c5948",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">10.0480</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">14.8652</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.0788</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.3516</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.6501</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.5272</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">12.2498</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">10.3057</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m10.0480\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m14.8652\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.0788\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.3516\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.6501\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.5272\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m12.2498\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m10.3057\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "A = E @ C \n",
    "LSE = torch.log( torch.sum( torch.exp(A), dim = 1, keepdim=True ) )\n",
    "print(LSE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b66582e4-85b8-4a6f-961e-42a1335f8ee6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m1\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">10.0480</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">14.8652</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.0788</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.3516</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.6501</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.5272</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">12.2498</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">10.3057</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m10.0480\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m14.8652\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.0788\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.3516\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.6501\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.5272\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m12.2498\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m10.3057\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CCE : Algorithm 2\n",
    "\n",
    "LSE = torch.ones(BS, 1) * -100000.0 # -inf\n",
    "\n",
    "# LSE part\n",
    "for n, e in enumerate(E_blocks):\n",
    "    for v, c in enumerate(C_blocks):\n",
    "        A = torch.zeros(n_b, v_b)\n",
    "        e_d_blocks = torch.split(e, d_b, dim = 1)\n",
    "        c_d_blocks = torch.split(c, d_b, dim = 0)\n",
    "        for ed, cd in zip(e_d_blocks, c_d_blocks):\n",
    "            A += ed @ cd\n",
    "        LSE_nv = torch.log( torch.sum( torch.exp(A), dim = 1, keepdim=True ) )\n",
    "        LSE[n_b * n : n_b * (n+1)] = torch.log( torch.exp(LSE[n_b * n : n_b * (n+1)]) \n",
    "                                                    + torch.exp(LSE_nv))\n",
    "print(LSE.shape)\n",
    "print(LSE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c96a6ee-64d1-4679-8a01-cdf116109ee3",
   "metadata": {},
   "source": [
    "#### Backward"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24289033-2754-4cd9-bad8-1577aaf294dc",
   "metadata": {},
   "source": [
    "#### standard backward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b008f356-49ca-4cd0-ace6-90b6bb6ff5bf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0914</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0.0914\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# basic backward\n",
    "# P = softmax(E @ C, dim = 1)\n",
    "# loss = - Y * torch.log(P)\n",
    "\n",
    "Y = torch.randn(BS, V)\n",
    "Y = torch.softmax(Y, dim = 1)\n",
    "\n",
    "X = E @ C\n",
    "P = torch.softmax(X, dim = 1)\n",
    "\n",
    "loss = (- Y * torch.log(P)).mean()\n",
    "print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f1e89020-c511-41ef-a5bb-ba7d6d02feaa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m128\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.0505e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5.9219e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.9728e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7721e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2872e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.4056e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.0426e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.4729e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.2635e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.3984e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.1512e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.5967e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.0976e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2592e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.8514e-04</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1033e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.8928e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.4670e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.1922e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3060e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0860e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2127e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.9582e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.4656e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.4197e-04</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-2.0505e-03\u001b[0m,  \u001b[1;36m5.9219e-05\u001b[0m, \u001b[1;36m-3.9728e-04\u001b[0m,  \u001b[1;36m1.7721e-03\u001b[0m, \u001b[1;36m-1.2872e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.4056e-03\u001b[0m, \u001b[1;36m-8.0426e-04\u001b[0m, \u001b[1;36m-1.4729e-03\u001b[0m,  \u001b[1;36m7.2635e-05\u001b[0m,  \u001b[1;36m1.3984e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m2.1512e-04\u001b[0m, \u001b[1;36m-6.5967e-05\u001b[0m, \u001b[1;36m-4.0976e-05\u001b[0m, \u001b[1;36m-1.2592e-04\u001b[0m, \u001b[1;36m-9.8514e-04\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-2.1033e-03\u001b[0m,  \u001b[1;36m7.8928e-04\u001b[0m, \u001b[1;36m-8.4670e-04\u001b[0m, \u001b[1;36m-6.1922e-04\u001b[0m, \u001b[1;36m-1.3060e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.0860e-03\u001b[0m, \u001b[1;36m-1.2127e-03\u001b[0m,  \u001b[1;36m1.9582e-05\u001b[0m,  \u001b[1;36m9.4656e-04\u001b[0m,  \u001b[1;36m1.4197e-04\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.0297e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.2597e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.6474e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7252e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.3637e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.8354e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.1001e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.7248e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1423e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.6895e-08</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.2884e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.9093e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.3355e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1693e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.0323e-06</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2662e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.1504e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.4147e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.0974e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0336e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.8036e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.4595e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.3332e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.7324e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.0484e-06</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-6.0297e-06\u001b[0m, \u001b[1;36m-9.2597e-06\u001b[0m,  \u001b[1;36m2.6474e-05\u001b[0m,  \u001b[1;36m1.7252e-05\u001b[0m,  \u001b[1;36m2.3637e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.8354e-07\u001b[0m, \u001b[1;36m-8.1001e-06\u001b[0m,  \u001b[1;36m2.7248e-05\u001b[0m, \u001b[1;36m-2.1423e-05\u001b[0m,  \u001b[1;36m6.6895e-08\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.2884e-06\u001b[0m,  \u001b[1;36m3.9093e-05\u001b[0m, \u001b[1;36m-5.3355e-05\u001b[0m,  \u001b[1;36m3.1693e-05\u001b[0m, \u001b[1;36m-9.0323e-06\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.2662e-05\u001b[0m,  \u001b[1;36m2.1504e-05\u001b[0m, \u001b[1;36m-8.4147e-05\u001b[0m, \u001b[1;36m-5.0974e-05\u001b[0m, \u001b[1;36m-1.0336e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.8036e-05\u001b[0m, \u001b[1;36m-2.4595e-05\u001b[0m, \u001b[1;36m-2.3332e-05\u001b[0m,  \u001b[1;36m3.7324e-06\u001b[0m, \u001b[1;36m-3.0484e-06\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Stand backward\n",
    "N = BS * V\n",
    "\n",
    "dP = (P - Y) / N\n",
    "dPdX = torch.ones(V, V)\n",
    "for i in range(BS):\n",
    "    dPdX += torch.diag(P[i,:]) - torch.outer(P[i,:], P[i,:]) \n",
    "\n",
    "dX = dP @ dPdX \n",
    "print(dX.shape)\n",
    "\n",
    "dE = dP @ C.t()\n",
    "dC = E.t() @ dP \n",
    "\n",
    "print(dE[:5,:5])\n",
    "print(dC[:5,:5])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cef7c1c-1d2f-4c7d-8e63-6bf39ad02c60",
   "metadata": {},
   "source": [
    "#### LSE Backward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "5744f1f9-418f-4731-a850-c3880b018399",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">10.0480</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">14.8652</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.0788</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.3516</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.6501</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.5272</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">12.2498</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">10.3057</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m10.0480\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m14.8652\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.0788\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.3516\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.6501\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.5272\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m12.2498\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m10.3057\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0914</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0.0914\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X = E @ C\n",
    "LSE = torch.log(torch.sum( torch.exp(X), dim = 1, keepdim = True))\n",
    "print(LSE)\n",
    "P = X - LSE # 逐行减去LSE\n",
    "loss = (- Y * P) # (-Y * X + Y * LSE)\n",
    "print(loss.mean() )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "6fd77d80-a146-410f-acd9-34ca19a5a937",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.0505e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5.9219e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.9728e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7721e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2872e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.4056e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.0426e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.4729e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.2635e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.3984e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.1512e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.5967e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.0976e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2592e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.8514e-04</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1033e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.8928e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.4670e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.1922e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3060e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0860e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2127e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.9582e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.4656e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.4197e-04</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-2.0505e-03\u001b[0m,  \u001b[1;36m5.9219e-05\u001b[0m, \u001b[1;36m-3.9728e-04\u001b[0m,  \u001b[1;36m1.7721e-03\u001b[0m, \u001b[1;36m-1.2872e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.4056e-03\u001b[0m, \u001b[1;36m-8.0426e-04\u001b[0m, \u001b[1;36m-1.4729e-03\u001b[0m,  \u001b[1;36m7.2635e-05\u001b[0m,  \u001b[1;36m1.3984e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m2.1512e-04\u001b[0m, \u001b[1;36m-6.5967e-05\u001b[0m, \u001b[1;36m-4.0976e-05\u001b[0m, \u001b[1;36m-1.2592e-04\u001b[0m, \u001b[1;36m-9.8514e-04\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-2.1033e-03\u001b[0m,  \u001b[1;36m7.8928e-04\u001b[0m, \u001b[1;36m-8.4670e-04\u001b[0m, \u001b[1;36m-6.1922e-04\u001b[0m, \u001b[1;36m-1.3060e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.0860e-03\u001b[0m, \u001b[1;36m-1.2127e-03\u001b[0m,  \u001b[1;36m1.9582e-05\u001b[0m,  \u001b[1;36m9.4656e-04\u001b[0m,  \u001b[1;36m1.4197e-04\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.0297e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.2597e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.6474e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7252e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.3637e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.8354e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.1001e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.7248e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1423e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.6895e-08</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.2884e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.9093e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.3355e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1693e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.0323e-06</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2662e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.1504e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.4147e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.0974e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0336e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.8036e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.4595e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.3332e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.7324e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.0484e-06</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-6.0297e-06\u001b[0m, \u001b[1;36m-9.2597e-06\u001b[0m,  \u001b[1;36m2.6474e-05\u001b[0m,  \u001b[1;36m1.7252e-05\u001b[0m,  \u001b[1;36m2.3637e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.8354e-07\u001b[0m, \u001b[1;36m-8.1001e-06\u001b[0m,  \u001b[1;36m2.7248e-05\u001b[0m, \u001b[1;36m-2.1423e-05\u001b[0m,  \u001b[1;36m6.6895e-08\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.2884e-06\u001b[0m,  \u001b[1;36m3.9093e-05\u001b[0m, \u001b[1;36m-5.3355e-05\u001b[0m,  \u001b[1;36m3.1693e-05\u001b[0m, \u001b[1;36m-9.0323e-06\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.2662e-05\u001b[0m,  \u001b[1;36m2.1504e-05\u001b[0m, \u001b[1;36m-8.4147e-05\u001b[0m, \u001b[1;36m-5.0974e-05\u001b[0m, \u001b[1;36m-1.0336e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.8036e-05\u001b[0m, \u001b[1;36m-2.4595e-05\u001b[0m, \u001b[1;36m-2.3332e-05\u001b[0m,  \u001b[1;36m3.7324e-06\u001b[0m, \u001b[1;36m-3.0484e-06\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dY_dX = -Y / N # P = \"X\" - LSE\n",
    "dY_dLSE = torch.sum(Y / N, dim = 1, keepdim = True) # P = X - \"LSE\"\n",
    "dLSE_dX = F.softmax(X, dim = 1)\n",
    "dY_dX_LSE = dY_dLSE * dLSE_dX\n",
    "\n",
    "dY_dX_total =  dY_dX + dY_dX_LSE\n",
    "\n",
    "dE = dY_dX_total @ C.t()\n",
    "dC = E.t() @ dY_dX_total \n",
    "\n",
    "print(dE[:5,:5])\n",
    "print(dC[:5,:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a0e5b051-06c7-4645-a200-51983b76b2b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1015</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0912</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.4023</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.6885</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3617</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.6390</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.9415</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3662</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0256</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.3215</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3928</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0740</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0271</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.2499</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0956</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1898</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9732</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.6191</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.6700</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3252</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.9124</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2397</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.0917</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7983</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0990</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-2.1015\u001b[0m,  \u001b[1;36m0.0912\u001b[0m, \u001b[1;36m-0.4023\u001b[0m,  \u001b[1;36m1.6885\u001b[0m, \u001b[1;36m-1.3617\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.6390\u001b[0m, \u001b[1;36m-0.9415\u001b[0m, \u001b[1;36m-1.3662\u001b[0m,  \u001b[1;36m0.0256\u001b[0m,  \u001b[1;36m1.3215\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.3928\u001b[0m,  \u001b[1;36m0.0740\u001b[0m, \u001b[1;36m-0.0271\u001b[0m, \u001b[1;36m-0.2499\u001b[0m, \u001b[1;36m-1.0956\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-2.1898\u001b[0m,  \u001b[1;36m0.9732\u001b[0m, \u001b[1;36m-0.6191\u001b[0m, \u001b[1;36m-0.6700\u001b[0m, \u001b[1;36m-1.3252\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-0.9124\u001b[0m, \u001b[1;36m-1.2397\u001b[0m, \u001b[1;36m-0.0917\u001b[0m,  \u001b[1;36m0.7983\u001b[0m,  \u001b[1;36m0.0990\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.5990e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.7586e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8.3423e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.6059e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.1319e-04</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.6817e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.5734e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.2944e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2376e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.7091e-04</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.8258e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.7603e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.6890e-02</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.5122e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.0148e-04</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2146e-02</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.9805e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.1456e-02</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.1811e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.5265e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.3457e-02</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.5087e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.2576e-02</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.9649e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.7716e-04</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-2.5990e-03\u001b[0m,  \u001b[1;36m3.7586e-05\u001b[0m,  \u001b[1;36m8.3423e-03\u001b[0m, \u001b[1;36m-1.6059e-03\u001b[0m, \u001b[1;36m-6.1319e-04\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m6.6817e-03\u001b[0m, \u001b[1;36m-5.5734e-05\u001b[0m, \u001b[1;36m-8.2944e-03\u001b[0m, \u001b[1;36m-1.2376e-04\u001b[0m, \u001b[1;36m-3.7091e-04\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-6.8258e-03\u001b[0m, \u001b[1;36m-1.7603e-04\u001b[0m, \u001b[1;36m-2.6890e-02\u001b[0m, \u001b[1;36m-8.5122e-04\u001b[0m, \u001b[1;36m-3.0148e-04\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.2146e-02\u001b[0m, \u001b[1;36m-2.9805e-04\u001b[0m, \u001b[1;36m-4.1456e-02\u001b[0m, \u001b[1;36m-3.1811e-05\u001b[0m, \u001b[1;36m-1.5265e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-3.3457e-02\u001b[0m, \u001b[1;36m-2.5087e-04\u001b[0m, \u001b[1;36m-4.2576e-02\u001b[0m, \u001b[1;36m-1.9649e-03\u001b[0m, \u001b[1;36m-2.7716e-04\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dE_LSE = dLSE_dX @ C.t()\n",
    "dC_LSE = E.t() @ dLSE_dX \n",
    "print(dE_LSE[:5,:5])\n",
    "print(dC_LSE[:5,:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0aafca2-b6c3-454e-8546-7256d10bcad7",
   "metadata": {},
   "source": [
    "#### Pytorch AutoGrad LSE backward"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9abc8949-c0de-430e-9823-1acd0387b670",
   "metadata": {},
   "source": [
    " 仅检验手动计算的梯度是否正确"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4b918e26-3eeb-4e19-b8a3-a142390e6337",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # 注意要清除梯度，得到的才准确\n",
    "# # E.grad.zero_()\n",
    "# # C.grad.zero_()\n",
    "# # Y.grad.zero_()\n",
    "\n",
    "# E.requires_grad = True\n",
    "# C.requires_grad = True\n",
    "# Y.requires_grad = True\n",
    "\n",
    "\n",
    "# X = E @ C\n",
    "# X.retain_grad()\n",
    "\n",
    "# LSE = torch.log(torch.sum( torch.exp(X), dim = 1, keepdim = True))\n",
    "# LSE.retain_grad()\n",
    "\n",
    "# P = X - LSE # 逐行减去LSE\n",
    "# P.retain_grad()\n",
    "\n",
    "# loss = (- Y * P) # (-Y * X + Y * LSE)\n",
    "# loss.retain_grad()\n",
    "# # print(loss)\n",
    "\n",
    "# l = loss.mean()\n",
    "# l.backward() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "f8c6e562-c15a-4c5b-808b-e3f63fb6bb77",
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(loss.grad[:5,:5])\n",
    "\n",
    "# print(LSE.grad[:5,:5])\n",
    "# print(LSE.grad.shape)\n",
    "\n",
    "# print(X.grad[:5,:5])\n",
    "# print(E.grad[:5,:5])\n",
    "# print(C.grad[:5,:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "210e91be-02e1-4787-9d95-98a462d5a47b",
   "metadata": {},
   "source": [
    "#### Algorithm 3: CCE LSE Backward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "ee70d263-ebff-46f3-b573-4da1c15372a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m128\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m16\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">16</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m16\u001b[0m, \u001b[1;36m128\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.0523e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8.9026e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.9288e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.6489e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3298e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.6006e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.1940e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3342e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.5002e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.2905e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.8360e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.2266e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.6462e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.4400e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0699e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1385e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.5039e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.0458e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.5433e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2941e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.9104e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2106e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.9572e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.7957e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.6675e-05</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-2.0523e-03\u001b[0m,  \u001b[1;36m8.9026e-05\u001b[0m, \u001b[1;36m-3.9288e-04\u001b[0m,  \u001b[1;36m1.6489e-03\u001b[0m, \u001b[1;36m-1.3298e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.6006e-03\u001b[0m, \u001b[1;36m-9.1940e-04\u001b[0m, \u001b[1;36m-1.3342e-03\u001b[0m,  \u001b[1;36m2.5002e-05\u001b[0m,  \u001b[1;36m1.2905e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.8360e-04\u001b[0m,  \u001b[1;36m7.2266e-05\u001b[0m, \u001b[1;36m-2.6462e-05\u001b[0m, \u001b[1;36m-2.4400e-04\u001b[0m, \u001b[1;36m-1.0699e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-2.1385e-03\u001b[0m,  \u001b[1;36m9.5039e-04\u001b[0m, \u001b[1;36m-6.0458e-04\u001b[0m, \u001b[1;36m-6.5433e-04\u001b[0m, \u001b[1;36m-1.2941e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-8.9104e-04\u001b[0m, \u001b[1;36m-1.2106e-03\u001b[0m, \u001b[1;36m-8.9572e-05\u001b[0m,  \u001b[1;36m7.7957e-04\u001b[0m,  \u001b[1;36m9.6675e-05\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.5381e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.6705e-08</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8.1468e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.5683e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.9882e-07</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.5251e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.4428e-08</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.1000e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2086e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.6222e-07</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.6658e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.7190e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.6260e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.3127e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.9442e-07</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1392e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.9107e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.0484e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.1066e-08</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.4908e-06</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.2672e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.4499e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.1578e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.9189e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.7066e-07</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-2.5381e-06\u001b[0m,  \u001b[1;36m3.6705e-08\u001b[0m,  \u001b[1;36m8.1468e-06\u001b[0m, \u001b[1;36m-1.5683e-06\u001b[0m, \u001b[1;36m-5.9882e-07\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m6.5251e-06\u001b[0m, \u001b[1;36m-5.4428e-08\u001b[0m, \u001b[1;36m-8.1000e-06\u001b[0m, \u001b[1;36m-1.2086e-07\u001b[0m, \u001b[1;36m-3.6222e-07\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-6.6658e-06\u001b[0m, \u001b[1;36m-1.7190e-07\u001b[0m, \u001b[1;36m-2.6260e-05\u001b[0m, \u001b[1;36m-8.3127e-07\u001b[0m, \u001b[1;36m-2.9442e-07\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.1392e-05\u001b[0m, \u001b[1;36m-2.9107e-07\u001b[0m, \u001b[1;36m-4.0484e-05\u001b[0m, \u001b[1;36m-3.1066e-08\u001b[0m, \u001b[1;36m-1.4908e-06\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-3.2672e-05\u001b[0m, \u001b[1;36m-2.4499e-07\u001b[0m, \u001b[1;36m-4.1578e-05\u001b[0m, \u001b[1;36m-1.9189e-06\u001b[0m, \u001b[1;36m-2.7066e-07\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CCE : Algorithm 3\n",
    "\n",
    "is_filter = False\n",
    "epsilon = 1e-6\n",
    "# \n",
    "dE = torch.zeros_like(E)\n",
    "dC = torch.zeros_like(C)\n",
    "\n",
    "dE_blocks = torch.split(dE, n_b, dim = 0)\n",
    "dC_blocks = torch.split(dC, v_b, dim = 1)\n",
    "\n",
    "print(Y.shape)\n",
    "d_LSE = dY_dLSE\n",
    "print(d_LSE)\n",
    "\n",
    "# LSE part\n",
    "for n, e in enumerate(E_blocks):\n",
    "    for v, c in enumerate(C_blocks):\n",
    "        \n",
    "        A = torch.zeros(n_b, v_b)\n",
    "        e_d_blocks = torch.split(e, d_b, dim = 1)\n",
    "        c_d_blocks = torch.split(c, d_b, dim = 0)\n",
    "        for ed, cd in zip(e_d_blocks, c_d_blocks):\n",
    "            A += ed @ cd # A_nv\n",
    "\n",
    "        S = torch.exp(A - LSE[n_b * n : n_b * (n+1)]) # softmax\n",
    "\n",
    "        # filter S\n",
    "        if is_filter:\n",
    "            if  (S < epsilon).all():\n",
    "                print('skip', n, v)\n",
    "                continue\n",
    "            \n",
    "        \n",
    "        d_LSE_block = d_LSE[n_b * n : n_b * (n+1), :]\n",
    "        d_index = 0\n",
    "        for ed, cd in zip(e_d_blocks, c_d_blocks):\n",
    "            dE_blocks[n][ :, d_index * d_b: (d_index+1)*d_b ] +=  (d_LSE_block * S) @ cd.t()\n",
    "            dC_blocks[v][ d_index * d_b: (d_index+1)*d_b, : ] +=  ed.t() @ (d_LSE_block * S)\n",
    "            d_index = d_index + 1\n",
    "\n",
    "\n",
    "# dLSE gradient part\n",
    "dE_part_LSE = torch.cat(dE_blocks, dim = 0)\n",
    "dC_part_LSE = torch.cat(dC_blocks, dim = 1)\n",
    "\n",
    "print(dE_part_LSE.shape)\n",
    "print(dC_part_LSE.shape)\n",
    "\n",
    "print(dE_part_LSE[:5,:5])\n",
    "print(dC_part_LSE[:5,:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "36854c45-3624-426b-9fb6-3390cd0ac420",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.0505e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5.9219e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.9728e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7721e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2872e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.4056e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.0426e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.4729e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.2635e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.3984e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.1512e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.5967e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.0976e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2592e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.8514e-04</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1033e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.8929e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.4670e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.1922e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3060e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0860e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2127e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.9582e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.4656e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.4197e-04</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-2.0505e-03\u001b[0m,  \u001b[1;36m5.9219e-05\u001b[0m, \u001b[1;36m-3.9728e-04\u001b[0m,  \u001b[1;36m1.7721e-03\u001b[0m, \u001b[1;36m-1.2872e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.4056e-03\u001b[0m, \u001b[1;36m-8.0426e-04\u001b[0m, \u001b[1;36m-1.4729e-03\u001b[0m,  \u001b[1;36m7.2635e-05\u001b[0m,  \u001b[1;36m1.3984e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m2.1512e-04\u001b[0m, \u001b[1;36m-6.5967e-05\u001b[0m, \u001b[1;36m-4.0976e-05\u001b[0m, \u001b[1;36m-1.2592e-04\u001b[0m, \u001b[1;36m-9.8514e-04\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-2.1033e-03\u001b[0m,  \u001b[1;36m7.8929e-04\u001b[0m, \u001b[1;36m-8.4670e-04\u001b[0m, \u001b[1;36m-6.1922e-04\u001b[0m, \u001b[1;36m-1.3060e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.0860e-03\u001b[0m, \u001b[1;36m-1.2127e-03\u001b[0m,  \u001b[1;36m1.9582e-05\u001b[0m,  \u001b[1;36m9.4656e-04\u001b[0m,  \u001b[1;36m1.4197e-04\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.0297e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.2597e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.6474e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7252e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.3637e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.8354e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.1001e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.7248e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1423e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.6895e-08</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.2884e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.9093e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.3355e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1693e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.0323e-06</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2662e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.1504e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.4147e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.0974e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0336e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.8036e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.4595e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.3332e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.7324e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.0484e-06</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-6.0297e-06\u001b[0m, \u001b[1;36m-9.2597e-06\u001b[0m,  \u001b[1;36m2.6474e-05\u001b[0m,  \u001b[1;36m1.7252e-05\u001b[0m,  \u001b[1;36m2.3637e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.8354e-07\u001b[0m, \u001b[1;36m-8.1001e-06\u001b[0m,  \u001b[1;36m2.7248e-05\u001b[0m, \u001b[1;36m-2.1423e-05\u001b[0m,  \u001b[1;36m6.6895e-08\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.2884e-06\u001b[0m,  \u001b[1;36m3.9093e-05\u001b[0m, \u001b[1;36m-5.3355e-05\u001b[0m,  \u001b[1;36m3.1693e-05\u001b[0m, \u001b[1;36m-9.0323e-06\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.2662e-05\u001b[0m,  \u001b[1;36m2.1504e-05\u001b[0m, \u001b[1;36m-8.4147e-05\u001b[0m, \u001b[1;36m-5.0974e-05\u001b[0m, \u001b[1;36m-1.0336e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.8036e-05\u001b[0m, \u001b[1;36m-2.4595e-05\u001b[0m, \u001b[1;36m-2.3332e-05\u001b[0m,  \u001b[1;36m3.7324e-06\u001b[0m, \u001b[1;36m-3.0484e-06\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# dX part dE, dC\n",
    "dE_part_X = dY_dX @ C.t()\n",
    "dC_part_X = E.t() @ dY_dX\n",
    "\n",
    "# dX Total\n",
    "dE_total = dE_part_X + dE_part_LSE\n",
    "dC_total = dC_part_X + dC_part_LSE\n",
    "\n",
    "print(dE_total[:5, :5])\n",
    "print(dC_total[:5, :5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67ea1a59-98f4-48a4-a8ac-823db6f57dce",
   "metadata": {},
   "source": [
    "#### Algorithm 4: CCE Backward total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "885e98ee-f301-4502-8e51-2e0630873ee8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([</span>  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">22</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">36</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">103</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">74</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">92</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">124</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">37</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m  \u001b[1;36m7\u001b[0m,  \u001b[1;36m22\u001b[0m,  \u001b[1;36m36\u001b[0m, \u001b[1;36m103\u001b[0m,  \u001b[1;36m74\u001b[0m,  \u001b[1;36m92\u001b[0m, \u001b[1;36m124\u001b[0m,  \u001b[1;36m37\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0611</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0.0611\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# print(Y)\n",
    "# Y_t = Y / 0.001 # be sharp\n",
    "# YY = F.softmax(Y_t, dim = 1)\n",
    "# print(YY)\n",
    "_, label = torch.max(Y, dim = 1)\n",
    "print(label)\n",
    "\n",
    "print(Y[0, label[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a041429-6553-4293-a859-b253fc9b4658",
   "metadata": {},
   "source": [
    "##### max likelihood version CE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "e5f78593-9655-494f-90cd-4307be24a95c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">([</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">128</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m8\u001b[0m, \u001b[1;36m128\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(P.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "ce6606b8-75a4-4bf2-bdaf-c878c7ceca6f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span>  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">22</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">36</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">103</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">74</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">92</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">124</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">37</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m  \u001b[1;36m7\u001b[0m,  \u001b[1;36m22\u001b[0m,  \u001b[1;36m36\u001b[0m, \u001b[1;36m103\u001b[0m,  \u001b[1;36m74\u001b[0m,  \u001b[1;36m92\u001b[0m, \u001b[1;36m124\u001b[0m,  \u001b[1;36m37\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4.2771</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">19.7241</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.0286</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">14.6713</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">12.1181</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">19.8909</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">18.0299</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">11.2865</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m \u001b[1;36m4.2771\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m19.7241\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.0286\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m14.6713\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m12.1181\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m19.8909\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m18.0299\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m11.2865\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">13.6283</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m13.6283\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(label.unsqueeze(0))\n",
    "P_CEL = torch.gather(P, dim = 1, index = label.unsqueeze(1))\n",
    "loss = -P_CEL\n",
    "l = loss.mean()\n",
    "print(loss)\n",
    "print(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "37e960f2-df94-461b-8564-fce9d6d12e16",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">13.6283</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m13.6283\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "loss = nn.CrossEntropyLoss()\n",
    "l_torch = loss(X, label)\n",
    "print(l_torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "c38e26f1-3009-4662-8aac-2007d53da5a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0010</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0010\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7035</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m1.7035\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(d_LSE)\n",
    "print(l_torch/BS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "c78677f7-083f-46da-a58a-a972635cec5a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.0505e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">5.9219e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.9728e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7721e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2872e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.4056e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.0426e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.4729e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.2635e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.3984e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.1512e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.5967e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-4.0976e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2592e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.8514e-04</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1033e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">7.8929e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.4670e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.1922e-04</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3060e-03</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0860e-03</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.2127e-03</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.9582e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.4656e-04</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.4197e-04</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-2.0505e-03\u001b[0m,  \u001b[1;36m5.9219e-05\u001b[0m, \u001b[1;36m-3.9728e-04\u001b[0m,  \u001b[1;36m1.7721e-03\u001b[0m, \u001b[1;36m-1.2872e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.4056e-03\u001b[0m, \u001b[1;36m-8.0426e-04\u001b[0m, \u001b[1;36m-1.4729e-03\u001b[0m,  \u001b[1;36m7.2635e-05\u001b[0m,  \u001b[1;36m1.3984e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m2.1512e-04\u001b[0m, \u001b[1;36m-6.5967e-05\u001b[0m, \u001b[1;36m-4.0976e-05\u001b[0m, \u001b[1;36m-1.2592e-04\u001b[0m, \u001b[1;36m-9.8514e-04\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-2.1033e-03\u001b[0m,  \u001b[1;36m7.8929e-04\u001b[0m, \u001b[1;36m-8.4670e-04\u001b[0m, \u001b[1;36m-6.1922e-04\u001b[0m, \u001b[1;36m-1.3060e-03\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.0860e-03\u001b[0m, \u001b[1;36m-1.2127e-03\u001b[0m,  \u001b[1;36m1.9582e-05\u001b[0m,  \u001b[1;36m9.4656e-04\u001b[0m,  \u001b[1;36m1.4197e-04\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">([[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-6.0297e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.2597e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.6474e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7252e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.3637e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.8354e-07</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.1001e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.7248e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.1423e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">6.6895e-08</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">9.2884e-06</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.9093e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.3355e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1693e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-9.0323e-06</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span> <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2662e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2.1504e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-8.4147e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-5.0974e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.0336e-05</span><span style=\"font-weight: bold\">]</span>,\n",
       "        <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.8036e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.4595e-05</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-2.3332e-05</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.7324e-06</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-3.0484e-06</span><span style=\"font-weight: bold\">]])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-6.0297e-06\u001b[0m, \u001b[1;36m-9.2597e-06\u001b[0m,  \u001b[1;36m2.6474e-05\u001b[0m,  \u001b[1;36m1.7252e-05\u001b[0m,  \u001b[1;36m2.3637e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.8354e-07\u001b[0m, \u001b[1;36m-8.1001e-06\u001b[0m,  \u001b[1;36m2.7248e-05\u001b[0m, \u001b[1;36m-2.1423e-05\u001b[0m,  \u001b[1;36m6.6895e-08\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m9.2884e-06\u001b[0m,  \u001b[1;36m3.9093e-05\u001b[0m, \u001b[1;36m-5.3355e-05\u001b[0m,  \u001b[1;36m3.1693e-05\u001b[0m, \u001b[1;36m-9.0323e-06\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.2662e-05\u001b[0m,  \u001b[1;36m2.1504e-05\u001b[0m, \u001b[1;36m-8.4147e-05\u001b[0m, \u001b[1;36m-5.0974e-05\u001b[0m, \u001b[1;36m-1.0336e-05\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.8036e-05\u001b[0m, \u001b[1;36m-2.4595e-05\u001b[0m, \u001b[1;36m-2.3332e-05\u001b[0m,  \u001b[1;36m3.7324e-06\u001b[0m, \u001b[1;36m-3.0484e-06\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CCE : Algorithm 4\n",
    "\n",
    "# M = torch.zeros(BS, V)\n",
    "# for i in range(BS):\n",
    "#     M[i, label[i]] = 1\n",
    "M_blocks = torch.split(Y, n_b, dim = 0)\n",
    "\n",
    "dE = torch.zeros_like(E)\n",
    "dC = torch.zeros_like(C)\n",
    "\n",
    "dE_blocks = torch.split(dE, n_b, dim = 0)\n",
    "dC_blocks = torch.split(dC, v_b, dim = 1)\n",
    "\n",
    "d_LCE = torch.ones(BS, 1) / (BS*V)\n",
    "\n",
    "\n",
    "for n, e in enumerate(E_blocks):\n",
    "    \n",
    "    M_v_blcoks = torch.split(M_blocks[n], v_b, dim = 1)\n",
    "                                 \n",
    "    for v, c in enumerate(C_blocks):\n",
    "        \n",
    "        A = torch.zeros(n_b, v_b)\n",
    "        e_d_blocks = torch.split(e, d_b, dim = 1)\n",
    "        c_d_blocks = torch.split(c, d_b, dim = 0)\n",
    "        for ed, cd in zip(e_d_blocks, c_d_blocks):\n",
    "            A += ed @ cd # A_nv\n",
    "\n",
    "        S = torch.exp(A - LSE[n_b * n : n_b * (n+1)]) \n",
    "        G = -M_v_blcoks[v] + S                          # (A-LSE) complete gradient\n",
    "                \n",
    "        # filter S\n",
    "        if is_filter:\n",
    "            if  (G.abs() < epsilon).all():\n",
    "                print('skip',n,v)\n",
    "                continue\n",
    "        \n",
    "        d_LCE_block = d_LCE[n_b * n : n_b * (n+1), :]\n",
    "        d_index = 0\n",
    "        for ed, cd in zip(e_d_blocks, c_d_blocks):\n",
    "            dE_blocks[n][ :, d_index * d_b: (d_index+1)*d_b ] +=  (d_LCE_block * G) @ cd.t()\n",
    "            dC_blocks[v][ d_index * d_b: (d_index+1)*d_b, : ] +=  ed.t() @ (d_LCE_block * G)\n",
    "            d_index = d_index + 1\n",
    "\n",
    "\n",
    "# total gradient for matrix E C\n",
    "dE = torch.cat(dE_blocks, dim = 0)\n",
    "dC = torch.cat(dC_blocks, dim = 1)\n",
    "\n",
    "print(dE[:5,:5])\n",
    "print(dC[:5,:5])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}