online_softmax_torch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6759a9b1-a71f-4f2d-b6f7-881d1c580261",
   "metadata": {},
   "source": [
    "# Online Softmax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9a81889-4e32-410f-9942-ae1e50197e73",
   "metadata": {},
   "source": [
    "github: xiaodongguaAIGC\n",
    "\n",
    "- softmax\n",
    "- Safe Softmax\n",
    "- online softmax\n",
    "- block online softmax\n",
    "- multi block online softmax\n",
    "- batch online softmax\n",
    "- multi block batch online softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ab5f591c-ac6d-4ddc-b7ef-b93d9e2187af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c7e9c2d6-63ff-4e55-ad3a-f8bec54c8fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.tensor([-0.3, 0.2, 0.5, 0.7, 0.1, 0.8])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e28e1fc-6bc2-4915-b2a5-04dc1b2b80dc",
   "metadata": {},
   "source": [
    "## Softmax By Torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bd32b22f-3dbe-45d2-ada1-74c13d3afe98",
   "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.0827</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1364</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1841</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2249</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1234</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2485</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.0827\u001b[0m, \u001b[1;36m0.1364\u001b[0m, \u001b[1;36m0.1841\u001b[0m, \u001b[1;36m0.2249\u001b[0m, \u001b[1;36m0.1234\u001b[0m, \u001b[1;36m0.2485\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_softmax = F.softmax(X, dim = 0)\n",
    "print(X_softmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0398b98-93fa-4a6d-b206-4e0fe3678c34",
   "metadata": {},
   "source": [
    "## Softmax By Handwrite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8761f59c-1954-4052-95fa-7e9cced802f4",
   "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.0827</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1364</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1841</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2249</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1234</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2485</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.0827\u001b[0m, \u001b[1;36m0.1364\u001b[0m, \u001b[1;36m0.1841\u001b[0m, \u001b[1;36m0.2249\u001b[0m, \u001b[1;36m0.1234\u001b[0m, \u001b[1;36m0.2485\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_exp_sum = X.exp().sum()\n",
    "X_softmax_hand = torch.exp(X) / X_exp_sum\n",
    "print(X_softmax_hand)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81247fc4-f14b-4ce2-a715-9f1200d63225",
   "metadata": {},
   "source": [
    "## Safe Softmax By Handwrite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "edba72aa-d000-4ae0-badc-1f9b171d9f7c",
   "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.0827</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1364</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1841</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2249</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1234</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2485</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.0827\u001b[0m, \u001b[1;36m0.1364\u001b[0m, \u001b[1;36m0.1841\u001b[0m, \u001b[1;36m0.2249\u001b[0m, \u001b[1;36m0.1234\u001b[0m, \u001b[1;36m0.2485\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_max = X.max()\n",
    "X_exp_sum_sub_max = torch.exp(X-X_max).sum()\n",
    "X_safe_softmax_hand = torch.exp(X - X_max) / X_exp_sum_sub_max\n",
    "print(X_safe_softmax_hand)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01273773-397c-4b0d-9607-01b1de1dea45",
   "metadata": {},
   "source": [
    "## Online Softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4f5f40da-622d-4745-b061-ccc392d32a7a",
   "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\">input x\n",
       "</pre>\n"
      ],
      "text/plain": [
       "input x\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.3000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.5000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.8000</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.3000\u001b[0m,  \u001b[1;36m0.2000\u001b[0m,  \u001b[1;36m0.5000\u001b[0m,  \u001b[1;36m0.7000\u001b[0m,  \u001b[1;36m0.1000\u001b[0m,  \u001b[1;36m0.8000\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.3000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.5000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1000</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.3000\u001b[0m,  \u001b[1;36m0.2000\u001b[0m,  \u001b[1;36m0.5000\u001b[0m,  \u001b[1;36m0.7000\u001b[0m,  \u001b[1;36m0.1000\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.8000</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m0.8000\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\">online softmax result:  <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.0827</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1364</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1841</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2249</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1234</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2485</span><span style=\"font-weight: bold\">])</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "online softmax result:  \u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0.0827\u001b[0m, \u001b[1;36m0.1364\u001b[0m, \u001b[1;36m0.1841\u001b[0m, \u001b[1;36m0.2249\u001b[0m, \u001b[1;36m0.1234\u001b[0m, \u001b[1;36m0.2485\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_pre = X[:-1]\n",
    "print('input x')\n",
    "print(X)\n",
    "print(X_pre)\n",
    "print(X[-1])\n",
    "\n",
    "# we calculative t-1 time Online Softmax\n",
    "X_max_pre = X_pre.max()\n",
    "X_sum_pre = torch.exp(X_pre - X_max_pre).sum()\n",
    "\n",
    "# we calculative t time Online Softmax\n",
    "X_max_cur = torch.max(X_max_pre, X[-1]) # X[-1] is new data\n",
    "X_sum_cur = X_sum_pre * torch.exp(X_max_pre - X_max_cur) + torch.exp(X[-1] - X_max_cur)\n",
    "\n",
    "# final we calculative online softmax\n",
    "X_online_softmax = torch.exp(X - X_max_cur) / X_sum_cur\n",
    "print('online softmax result: ', X_online_softmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b80860f-c775-4992-a580-a5c91e080670",
   "metadata": {},
   "source": [
    "## Block Online Softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c6ac9ff3-e6df-4cd9-8fd6-2d491e41da23",
   "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.3000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.5000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.8000</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.3000\u001b[0m,  \u001b[1;36m0.2000\u001b[0m,  \u001b[1;36m0.5000\u001b[0m,  \u001b[1;36m0.7000\u001b[0m,  \u001b[1;36m0.1000\u001b[0m,  \u001b[1;36m0.8000\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=\"font-weight: bold\">(</span><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.3000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.5000</span><span style=\"font-weight: bold\">])</span>, <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.7000</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1000</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.8000</span><span style=\"font-weight: bold\">]))</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m(\u001b[0m\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-0.3000\u001b[0m,  \u001b[1;36m0.2000\u001b[0m,  \u001b[1;36m0.5000\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m, \u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0.7000\u001b[0m, \u001b[1;36m0.1000\u001b[0m, \u001b[1;36m0.8000\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_block = torch.split(X, split_size_or_sections = 3 , dim = 0) \n",
    "print(X)\n",
    "print(X_block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "83c30afc-84ac-4605-8a7f-f8360a6a3c2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we parallel calculate  different block max & sum\n",
    "X_block_0_max = X_block[0].max()\n",
    "X_block_0_sum = torch.exp(X_block[0] - X_block_0_max).sum()\n",
    "\n",
    "X_block_1_max = X_block[1].max()\n",
    "X_block_1_sum = torch.exp(X_block[1] - X_block_1_max).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f4e18507-5238-456e-956f-06568376b173",
   "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.0827</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1364</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1841</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2249</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1234</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2485</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.0827\u001b[0m, \u001b[1;36m0.1364\u001b[0m, \u001b[1;36m0.1841\u001b[0m, \u001b[1;36m0.2249\u001b[0m, \u001b[1;36m0.1234\u001b[0m, \u001b[1;36m0.2485\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# parallel online block update max & sum\n",
    "X_max_global = torch.max(X_block_0_max, X_block_1_max) \n",
    "L_global = (X_block_0_sum * torch.exp(X_block_0_max - X_max_global) \\\n",
    "            + X_block_1_sum * torch.exp(X_block_1_max - X_max_global)) # block sum\n",
    "\n",
    "X_block_online_softmax_parallel = torch.exp(X - X_max_global) / L_global\n",
    "print(X_block_online_softmax_parallel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9f618de1-ec3d-4d14-a10b-77e291480249",
   "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.0827</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1364</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1841</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2249</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1234</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2485</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.0827\u001b[0m, \u001b[1;36m0.1364\u001b[0m, \u001b[1;36m0.1841\u001b[0m, \u001b[1;36m0.2249\u001b[0m, \u001b[1;36m0.1234\u001b[0m, \u001b[1;36m0.2485\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# online block update max & sum\n",
    "# updated version for multi-block, simpler version\n",
    "X_block_1_max_update = torch.max(X_block_0_max, X_block_1_max) \n",
    "X_block_1_sum_update = X_block_0_sum * torch.exp(X_block_0_max - X_block_1_max_update) \\\n",
    "                     + torch.exp(X_block[1] - X_block_1_max_update).sum() # block sum\n",
    "\n",
    "X_block_online_softmax = torch.exp(X - X_block_1_max_update) / X_block_1_sum_update\n",
    "print(X_block_online_softmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d68937aa-1d6d-4f3a-a013-0ef6fb763607",
   "metadata": {},
   "source": [
    "## Multi Block Online Softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7a815125-e36d-4098-924c-abfbdac07fd4",
   "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.3000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.5000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.8000</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.3000\u001b[0m,  \u001b[1;36m0.2000\u001b[0m,  \u001b[1;36m0.5000\u001b[0m,  \u001b[1;36m0.7000\u001b[0m,  \u001b[1;36m0.1000\u001b[0m,  \u001b[1;36m0.8000\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=\"font-weight: bold\">(</span><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.3000</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2000</span><span style=\"font-weight: bold\">])</span>, <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.5000</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7000</span><span style=\"font-weight: bold\">])</span>, <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.1000</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.8000</span><span style=\"font-weight: bold\">]))</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m(\u001b[0m\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-0.3000\u001b[0m,  \u001b[1;36m0.2000\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m, \u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0.5000\u001b[0m, \u001b[1;36m0.7000\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m, \u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m0.1000\u001b[0m, \u001b[1;36m0.8000\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_block = torch.split(X, split_size_or_sections = 2, dim = 0) \n",
    "print(X)\n",
    "print(X_block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "26ed247b-2d4b-478a-a9d3-328cb23e0073",
   "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.0827</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1364</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1841</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2249</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1234</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2485</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.0827\u001b[0m, \u001b[1;36m0.1364\u001b[0m, \u001b[1;36m0.1841\u001b[0m, \u001b[1;36m0.2249\u001b[0m, \u001b[1;36m0.1234\u001b[0m, \u001b[1;36m0.2485\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.0000</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m1.0000\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# online multi-block update max & sum\n",
    "M_old = torch.tensor([-100000.0])\n",
    "L_old = torch.tensor([0.0])\n",
    "\n",
    "for i in range(len(X_block)):\n",
    "    M = torch.max(X_block[i])\n",
    "    M_new = torch.max(M, M_old) \n",
    "    \n",
    "    L_new = L_old * torch.exp(M_old - M_new) \\\n",
    "            +  torch.exp(X_block[i] - M).sum() * torch.exp(M - M_new) \n",
    "    \n",
    "    # use simplest format\n",
    "    # L_new = L_old * torch.exp(M_old - M_new) \\\n",
    "    #         +  torch.exp(X_block[i] - M_new).sum() \n",
    "    \n",
    "    M_old = M_new\n",
    "    L_old = L_new\n",
    "\n",
    "X_multi_block_online_softmax = torch.exp(X - M_old) / L_old\n",
    "print(X_multi_block_online_softmax)\n",
    "print(X_multi_block_online_softmax.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b0aadfb-df9a-4835-9945-b7c5b40bd995",
   "metadata": {},
   "source": [
    "## Batch Online Softmax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c5288f-495e-415e-8a01-d9a683b489ef",
   "metadata": {},
   "source": [
    "### Batch Online Softmax by Torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "57e6e18e-4098-439e-8d97-094041d48f90",
   "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\">1.2568</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7383</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3952</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0712</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.4927</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.4833</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.0294</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.8888</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1123</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7380</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3930</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9524</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.3562</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.3688</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.3282</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2616</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.5259</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1031</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.1950</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3002</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.2051</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3824</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.6178</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9580</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;36m1.2568\u001b[0m,  \u001b[1;36m0.7383\u001b[0m,  \u001b[1;36m0.3952\u001b[0m,  \u001b[1;36m1.0712\u001b[0m, \u001b[1;36m-0.4927\u001b[0m, \u001b[1;36m-1.4833\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.0294\u001b[0m,  \u001b[1;36m0.8888\u001b[0m, \u001b[1;36m-0.1123\u001b[0m,  \u001b[1;36m1.7380\u001b[0m, \u001b[1;36m-1.3930\u001b[0m,  \u001b[1;36m0.9524\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.3562\u001b[0m, \u001b[1;36m-0.3688\u001b[0m, \u001b[1;36m-0.3282\u001b[0m,  \u001b[1;36m3.2616\u001b[0m,  \u001b[1;36m1.5259\u001b[0m, \u001b[1;36m-0.1031\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.1950\u001b[0m,  \u001b[1;36m0.3002\u001b[0m,  \u001b[1;36m1.2051\u001b[0m, \u001b[1;36m-1.3824\u001b[0m,  \u001b[1;36m1.6178\u001b[0m,  \u001b[1;36m0.9580\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.3240</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1929</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1369</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2691</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0563</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0209</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.0799</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1888</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0694</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.4414</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0193</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2012</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.0415</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0201</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0209</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7578</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1336</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0262</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.0881</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0978</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2418</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0182</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3653</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1888</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.3240\u001b[0m, \u001b[1;36m0.1929\u001b[0m, \u001b[1;36m0.1369\u001b[0m, \u001b[1;36m0.2691\u001b[0m, \u001b[1;36m0.0563\u001b[0m, \u001b[1;36m0.0209\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0799\u001b[0m, \u001b[1;36m0.1888\u001b[0m, \u001b[1;36m0.0694\u001b[0m, \u001b[1;36m0.4414\u001b[0m, \u001b[1;36m0.0193\u001b[0m, \u001b[1;36m0.2012\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0415\u001b[0m, \u001b[1;36m0.0201\u001b[0m, \u001b[1;36m0.0209\u001b[0m, \u001b[1;36m0.7578\u001b[0m, \u001b[1;36m0.1336\u001b[0m, \u001b[1;36m0.0262\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0881\u001b[0m, \u001b[1;36m0.0978\u001b[0m, \u001b[1;36m0.2418\u001b[0m, \u001b[1;36m0.0182\u001b[0m, \u001b[1;36m0.3653\u001b[0m, \u001b[1;36m0.1888\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.0000</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0000</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0000</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0000</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;36m1.0000\u001b[0m, \u001b[1;36m1.0000\u001b[0m, \u001b[1;36m1.0000\u001b[0m, \u001b[1;36m1.0000\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_batch = torch.randn(4, 6)\n",
    "print(X_batch)\n",
    "X_batch_softmax = F.softmax(X_batch, dim = 1) \n",
    "print(X_batch_softmax)\n",
    "X_batch_softmax_evaluete = X_batch_softmax.sum(dim = 1)\n",
    "print(X_batch_softmax_evaluete) # row prob sum is 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9821e291-a47a-48b8-9c63-2f39d4960867",
   "metadata": {},
   "source": [
    "### Batch Online Softmax by Hand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b8abed75-6800-4cdb-be7b-a3a958a32256",
   "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\">3</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;36m4\u001b[0m \u001b[1;36m3\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.2568</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7383</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3952</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0712</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.4927</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.4833</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.0294</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.8888</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1123</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7380</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3930</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9524</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.3562</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.3688</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.3282</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2616</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.5259</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1031</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.1950</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3002</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.2051</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3824</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.6178</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9580</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;36m1.2568\u001b[0m,  \u001b[1;36m0.7383\u001b[0m,  \u001b[1;36m0.3952\u001b[0m,  \u001b[1;36m1.0712\u001b[0m, \u001b[1;36m-0.4927\u001b[0m, \u001b[1;36m-1.4833\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.0294\u001b[0m,  \u001b[1;36m0.8888\u001b[0m, \u001b[1;36m-0.1123\u001b[0m,  \u001b[1;36m1.7380\u001b[0m, \u001b[1;36m-1.3930\u001b[0m,  \u001b[1;36m0.9524\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.3562\u001b[0m, \u001b[1;36m-0.3688\u001b[0m, \u001b[1;36m-0.3282\u001b[0m,  \u001b[1;36m3.2616\u001b[0m,  \u001b[1;36m1.5259\u001b[0m, \u001b[1;36m-0.1031\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.1950\u001b[0m,  \u001b[1;36m0.3002\u001b[0m,  \u001b[1;36m1.2051\u001b[0m, \u001b[1;36m-1.3824\u001b[0m,  \u001b[1;36m1.6178\u001b[0m,  \u001b[1;36m0.9580\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.2568</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7383</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3952</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.0294</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.8888</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1123</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.3562</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.3688</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.3282</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.1950</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3002</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.2051</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;36m1.2568\u001b[0m,  \u001b[1;36m0.7383\u001b[0m,  \u001b[1;36m0.3952\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.0294\u001b[0m,  \u001b[1;36m0.8888\u001b[0m, \u001b[1;36m-0.1123\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.3562\u001b[0m, \u001b[1;36m-0.3688\u001b[0m, \u001b[1;36m-0.3282\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.1950\u001b[0m,  \u001b[1;36m0.3002\u001b[0m,  \u001b[1;36m1.2051\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.0712</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.4927</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.4833</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.7380</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3930</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9524</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.2616</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.5259</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1031</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.3824</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.6178</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9580</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;36m1.0712\u001b[0m, \u001b[1;36m-0.4927\u001b[0m, \u001b[1;36m-1.4833\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.7380\u001b[0m, \u001b[1;36m-1.3930\u001b[0m,  \u001b[1;36m0.9524\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m3.2616\u001b[0m,  \u001b[1;36m1.5259\u001b[0m, \u001b[1;36m-0.1031\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.3824\u001b[0m,  \u001b[1;36m1.6178\u001b[0m,  \u001b[1;36m0.9580\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "b, d = X_batch.shape\n",
    "print(b, d//2)\n",
    "\n",
    "X_batch_block_0 = X_batch[:, :d//2]\n",
    "X_batch_block_1 = X_batch[:, d//2:]\n",
    "\n",
    "print(X_batch)\n",
    "print(X_batch_block_0)\n",
    "print(X_batch_block_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8c89f06f-5536-42ef-880e-55359005c1a5",
   "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\">1.2568</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.8888</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.3562</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.2051</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;36m1.2568\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.8888\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.3562\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m1.2051\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.0179</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.7909</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.9887</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.7688</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.0179\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m1.7909\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m1.9887\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m1.7688\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# we parallel calculate  different block max & sum\n",
    "X_batch_0_max, _ = X_batch_block_0.max(dim = 1, keepdim = True)\n",
    "X_batch_0_sum = torch.exp(X_batch_block_0 - X_batch_0_max).sum(dim = 1, keepdim = True)\n",
    "\n",
    "X_batch_1_max, _ = X_batch_block_1.max(dim = 1, keepdim = True)\n",
    "X_batch_1_sum = torch.exp(X_batch_block_1 - X_batch_1_max).sum(dim = 1, keepdim = True)\n",
    "\n",
    "print(X_batch_0_max)\n",
    "print(X_batch_0_sum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9b97605a-4d38-40d1-bb90-c4243c608daa",
   "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.3240</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1929</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1369</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2691</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0563</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0209</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.0799</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1888</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0694</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.4414</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0193</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2012</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.0415</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0201</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0209</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7578</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1336</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0262</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.0881</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0978</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2418</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0182</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3653</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1888</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.3240\u001b[0m, \u001b[1;36m0.1929\u001b[0m, \u001b[1;36m0.1369\u001b[0m, \u001b[1;36m0.2691\u001b[0m, \u001b[1;36m0.0563\u001b[0m, \u001b[1;36m0.0209\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0799\u001b[0m, \u001b[1;36m0.1888\u001b[0m, \u001b[1;36m0.0694\u001b[0m, \u001b[1;36m0.4414\u001b[0m, \u001b[1;36m0.0193\u001b[0m, \u001b[1;36m0.2012\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0415\u001b[0m, \u001b[1;36m0.0201\u001b[0m, \u001b[1;36m0.0209\u001b[0m, \u001b[1;36m0.7578\u001b[0m, \u001b[1;36m0.1336\u001b[0m, \u001b[1;36m0.0262\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0881\u001b[0m, \u001b[1;36m0.0978\u001b[0m, \u001b[1;36m0.2418\u001b[0m, \u001b[1;36m0.0182\u001b[0m, \u001b[1;36m0.3653\u001b[0m, \u001b[1;36m0.1888\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# online batch block update max & sum\n",
    "X_batch_1_max_update = torch.maximum(X_batch_0_max, X_batch_1_max) # 逐个元素找最大值\n",
    "X_batch_1_sum_update = X_batch_0_sum * torch.exp(X_batch_0_max - X_batch_1_max_update) \\\n",
    "                     + torch.exp(X_batch_block_1 - X_batch_1_max_update).sum(dim = 1, keepdim = True) # block sum\n",
    "\n",
    "X_batch_online_softmax = torch.exp(X_batch - X_batch_1_max_update) / X_batch_1_sum_update\n",
    "print(X_batch_online_softmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "540be634-264d-454f-a242-18cd0680ac38",
   "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.3240</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1929</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1369</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2691</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0563</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0209</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.0799</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1888</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0694</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.4414</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0193</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2012</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.0415</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0201</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0209</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7578</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1336</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0262</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.0881</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0978</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2418</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0182</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3653</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1888</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.3240\u001b[0m, \u001b[1;36m0.1929\u001b[0m, \u001b[1;36m0.1369\u001b[0m, \u001b[1;36m0.2691\u001b[0m, \u001b[1;36m0.0563\u001b[0m, \u001b[1;36m0.0209\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0799\u001b[0m, \u001b[1;36m0.1888\u001b[0m, \u001b[1;36m0.0694\u001b[0m, \u001b[1;36m0.4414\u001b[0m, \u001b[1;36m0.0193\u001b[0m, \u001b[1;36m0.2012\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0415\u001b[0m, \u001b[1;36m0.0201\u001b[0m, \u001b[1;36m0.0209\u001b[0m, \u001b[1;36m0.7578\u001b[0m, \u001b[1;36m0.1336\u001b[0m, \u001b[1;36m0.0262\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0881\u001b[0m, \u001b[1;36m0.0978\u001b[0m, \u001b[1;36m0.2418\u001b[0m, \u001b[1;36m0.0182\u001b[0m, \u001b[1;36m0.3653\u001b[0m, \u001b[1;36m0.1888\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_batch_softmax_torch = F.softmax(X_batch, dim = 1) \n",
    "print(X_batch_softmax_torch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87ec5c37-f50d-44bc-8643-14aeb69018e1",
   "metadata": {},
   "source": [
    "### Multi Block Batch Online Softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "19f1d55f-a943-460e-8fcf-04798b1d4423",
   "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=\"font-weight: bold\">(</span>\n",
       "    <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.2568</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7383</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.0294</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.8888</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.3562</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.3688</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.1950</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3002</span><span style=\"font-weight: bold\">]])</span>,\n",
       "    <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.3952</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.0712</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.1123</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1.7380</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.3282</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2616</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.2051</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.3824</span><span style=\"font-weight: bold\">]])</span>,\n",
       "    <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.4927</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-1.4833</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.3930</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9524</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.5259</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">-0.1031</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.6178</span>,  <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.9580</span><span style=\"font-weight: bold\">]])</span>\n",
       "<span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m(\u001b[0m\n",
       "    \u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m \u001b[1;36m1.2568\u001b[0m,  \u001b[1;36m0.7383\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.0294\u001b[0m,  \u001b[1;36m0.8888\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.3562\u001b[0m, \u001b[1;36m-0.3688\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m0.1950\u001b[0m,  \u001b[1;36m0.3002\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m,\n",
       "    \u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m \u001b[1;36m0.3952\u001b[0m,  \u001b[1;36m1.0712\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-0.1123\u001b[0m,  \u001b[1;36m1.7380\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-0.3282\u001b[0m,  \u001b[1;36m3.2616\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.2051\u001b[0m, \u001b[1;36m-1.3824\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m,\n",
       "    \u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[1m[\u001b[0m\u001b[1m[\u001b[0m\u001b[1;36m-0.4927\u001b[0m, \u001b[1;36m-1.4833\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m-1.3930\u001b[0m,  \u001b[1;36m0.9524\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.5259\u001b[0m, \u001b[1;36m-0.1031\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m \u001b[1;36m1.6178\u001b[0m,  \u001b[1;36m0.9580\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n",
       "\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# X_batch = torch.randn(4, 6)\n",
    "X_blocks = torch.split(X_batch, 2, dim=1)\n",
    "print(X_blocks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c2ff29eb-9d6b-43ae-b7af-d6bf6c594b5d",
   "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.3240</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1929</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1369</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2691</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0563</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0209</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.0799</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1888</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0694</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.4414</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0193</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2012</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.0415</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0201</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0209</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.7578</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1336</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0262</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.0881</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0978</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.2418</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.0182</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.3653</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1888</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.3240\u001b[0m, \u001b[1;36m0.1929\u001b[0m, \u001b[1;36m0.1369\u001b[0m, \u001b[1;36m0.2691\u001b[0m, \u001b[1;36m0.0563\u001b[0m, \u001b[1;36m0.0209\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0799\u001b[0m, \u001b[1;36m0.1888\u001b[0m, \u001b[1;36m0.0694\u001b[0m, \u001b[1;36m0.4414\u001b[0m, \u001b[1;36m0.0193\u001b[0m, \u001b[1;36m0.2012\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0415\u001b[0m, \u001b[1;36m0.0201\u001b[0m, \u001b[1;36m0.0209\u001b[0m, \u001b[1;36m0.7578\u001b[0m, \u001b[1;36m0.1336\u001b[0m, \u001b[1;36m0.0262\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m0.0881\u001b[0m, \u001b[1;36m0.0978\u001b[0m, \u001b[1;36m0.2418\u001b[0m, \u001b[1;36m0.0182\u001b[0m, \u001b[1;36m0.3653\u001b[0m, \u001b[1;36m0.1888\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.0000</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.0000</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.0000</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.0000</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;36m1.0000\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m1.0000\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m1.0000\u001b[0m\u001b[1m]\u001b[0m,\n",
       "        \u001b[1m[\u001b[0m\u001b[1;36m1.0000\u001b[0m\u001b[1m]\u001b[0m\u001b[1m]\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "b, d = X_batch.shape\n",
    "M_old = torch.ones((b,1)) * -100000.0\n",
    "L_old = torch.zeros((b,1))\n",
    "\n",
    "for X_block in X_blocks:\n",
    "    M,_ = torch.max(X_block, dim = 1, keepdim = True)\n",
    "    M_new = torch.maximum(M, M_old) \n",
    "    \n",
    "    L_new = L_old * torch.exp(M_old - M_new) \\\n",
    "            + torch.exp(X_block - M_new).sum(dim = 1, keepdim = True) \n",
    "    \n",
    "    M_old = M_new\n",
    "    L_old = L_new\n",
    "\n",
    "X_blocks_batch = torch.exp(X_batch - M_old) / L_old\n",
    "print(X_blocks_batch)\n",
    "print(X_blocks_batch.sum(dim = 1, keepdim = True))"
   ]
  }
 ],
 "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
}