{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vmap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## vmap()\n",
    "\n",
    "vmap is another of the key components of jax. What it does is that it lets you pass a batch of inputs to your functions in a single go instead of passing them one by one. Think of it like this. You have 10 arrays which you need to pass to a function. In a typical python program, you would be writing a loop. With vmap, you can send them as a batch, take advantage of vector and matrix operations supported by your accelerator (GPU, TPU etc.) and make your computation way faster. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's consider a simple affine transformation over some arrays.\n",
    "\n",
    "$$\n",
    "y = xW^{T} + b\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:46:45.179981Z",
     "start_time": "2023-02-08T02:46:44.579860Z"
    }
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "key = jax.random.key(42)\n",
    "key, *subkeys = jax.random.split(key, 3) # main key + 2 keys more"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:46:45.183481Z",
     "start_time": "2023-02-08T02:46:45.181629Z"
    }
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def affine(x, w, b):\n",
    "    return x @ w.T + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:46:45.478617Z",
     "start_time": "2023-02-08T02:46:45.184566Z"
    }
   },
   "outputs": [],
   "source": [
    "# say for this affine transformation, we are taking \n",
    "# a 10 dim vector\n",
    "# and making it a 2 dim vector\n",
    "\n",
    "in_dim = 10\n",
    "out_dim = 2\n",
    "\n",
    "x = jax.random.normal(key, shape=(in_dim, ))\n",
    "\n",
    "w = jax.random.normal(subkeys[0], shape=(out_dim, in_dim))\n",
    "\n",
    "b = jax.random.normal(subkeys[1], shape=(out_dim, ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:46:45.522829Z",
     "start_time": "2023-02-08T02:46:45.479643Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([-0.98225653,  1.9490494 ], dtype=float32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "affine(x, w, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But that was over one `x`, `w` and `b`. What if we want to pass a batch as intended?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:46:45.791286Z",
     "start_time": "2023-02-08T02:46:45.523792Z"
    }
   },
   "outputs": [],
   "source": [
    "batch_size = 10\n",
    "\n",
    "xs = jax.random.normal(key, shape=(batch_size, 10, ))\n",
    "ws = jax.random.normal(subkeys[0], shape=(batch_size, out_dim, in_dim))\n",
    "bs = jax.random.normal(subkeys[1], shape=(batch_size, out_dim, ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to add $b$ and make it a full fledged affine transformation .... "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:46:47.178875Z",
     "start_time": "2023-02-08T02:46:45.792843Z"
    }
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "dot_general requires contracting dimensions to have the same shape, got (10,) and (2,).",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mTypeError\u001b[39m                                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[7]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43maffine\u001b[49m\u001b[43m(\u001b[49m\u001b[43mxs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mws\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbs\u001b[49m\u001b[43m)\u001b[49m\n",
      "    \u001b[31m[... skipping hidden 13 frame]\u001b[39m\n",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 3\u001b[39m, in \u001b[36maffine\u001b[39m\u001b[34m(x, w, b)\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;129m@jax\u001b[39m.jit\n\u001b[32m      2\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34maffine\u001b[39m(x, w, b):\n\u001b[32m----> \u001b[39m\u001b[32m3\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[43m@\u001b[49m\u001b[43m \u001b[49m\u001b[43mw\u001b[49m\u001b[43m.\u001b[49m\u001b[43mT\u001b[49m + b\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:1141\u001b[39m, in \u001b[36m_forward_operator_to_aval.<locals>.op\u001b[39m\u001b[34m(self, *args)\u001b[39m\n\u001b[32m   1140\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mop\u001b[39m(\u001b[38;5;28mself\u001b[39m, *args):\n\u001b[32m-> \u001b[39m\u001b[32m1141\u001b[39m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mgetattr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43maval\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[33;43mf\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43m_\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mname\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:604\u001b[39m, in \u001b[36m_defer_to_unrecognized_arg.<locals>.deferring_binary_op\u001b[39m\u001b[34m(self, other)\u001b[39m\n\u001b[32m    602\u001b[39m args = (other, \u001b[38;5;28mself\u001b[39m) \u001b[38;5;28;01mif\u001b[39;00m swap \u001b[38;5;28;01melse\u001b[39;00m (\u001b[38;5;28mself\u001b[39m, other)\n\u001b[32m    603\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(other, _accepted_binop_types):\n\u001b[32m--> \u001b[39m\u001b[32m604\u001b[39m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mbinary_op\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    605\u001b[39m \u001b[38;5;66;03m# Note: don't use isinstance here, because we don't want to raise for\u001b[39;00m\n\u001b[32m    606\u001b[39m \u001b[38;5;66;03m# subclasses, e.g. NamedTuple objects that may override operators.\u001b[39;00m\n\u001b[32m    607\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(other) \u001b[38;5;129;01min\u001b[39;00m _rejected_binop_types:\n",
      "    \u001b[31m[... skipping hidden 14 frame]\u001b[39m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/tensor_contractions.py:254\u001b[39m, in \u001b[36mmatmul\u001b[39m\u001b[34m(a, b, precision, preferred_element_type, out_sharding)\u001b[39m\n\u001b[32m    252\u001b[39m a = lax.squeeze(a, \u001b[38;5;28mtuple\u001b[39m(a_squeeze))\n\u001b[32m    253\u001b[39m b = lax.squeeze(b, \u001b[38;5;28mtuple\u001b[39m(b_squeeze))\n\u001b[32m--> \u001b[39m\u001b[32m254\u001b[39m out = \u001b[43mlax\u001b[49m\u001b[43m.\u001b[49m\u001b[43mdot_general\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    255\u001b[39m \u001b[43m  \u001b[49m\u001b[43ma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mb\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnp\u001b[49m\u001b[43m.\u001b[49m\u001b[43mndim\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[43m-\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[43mnp\u001b[49m\u001b[43m.\u001b[49m\u001b[43mndim\u001b[49m\u001b[43m(\u001b[49m\u001b[43mb\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[43m-\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m1\u001b[39;49m\u001b[43m \u001b[49m\u001b[43m-\u001b[49m\u001b[43m \u001b[49m\u001b[43mb_is_mat\u001b[49m\u001b[43m,\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m(\u001b[49m\u001b[43ma_batch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mb_batch\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    256\u001b[39m \u001b[43m  \u001b[49m\u001b[43mprecision\u001b[49m\u001b[43m=\u001b[49m\u001b[43mprecision\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpreferred_element_type\u001b[49m\u001b[43m=\u001b[49m\u001b[43mpreferred_element_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    257\u001b[39m \u001b[43m  \u001b[49m\u001b[43mout_sharding\u001b[49m\u001b[43m=\u001b[49m\u001b[43mout_sharding\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    258\u001b[39m result = lax.transpose(out, perm)\n\u001b[32m    259\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m lax._convert_element_type(result, preferred_element_type, output_weak_type)\n",
      "    \u001b[31m[... skipping hidden 13 frame]\u001b[39m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/lax/lax.py:5303\u001b[39m, in \u001b[36m_dot_general_shape_rule\u001b[39m\u001b[34m(lhs, rhs, dimension_numbers, precision, preferred_element_type, out_sharding)\u001b[39m\n\u001b[32m   5300\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m core.definitely_equal_shape(lhs_contracting_shape, rhs_contracting_shape):\n\u001b[32m   5301\u001b[39m   msg = (\u001b[33m\"\u001b[39m\u001b[33mdot_general requires contracting dimensions to have the same \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m   5302\u001b[39m          \u001b[33m\"\u001b[39m\u001b[33mshape, got \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[33m and \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[33m.\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m-> \u001b[39m\u001b[32m5303\u001b[39m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(msg.format(lhs_contracting_shape, rhs_contracting_shape))\n\u001b[32m   5305\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m _dot_general_shape_computation(lhs.shape, rhs.shape, dimension_numbers)\n",
      "\u001b[31mTypeError\u001b[39m: dot_general requires contracting dimensions to have the same shape, got (10,) and (2,)."
     ]
    }
   ],
   "source": [
    "# won't run due to shape mismatch\n",
    "affine(xs, ws, bs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:27:33.821783Z",
     "start_time": "2023-02-08T02:27:33.819156Z"
    }
   },
   "source": [
    "## Enter vmap:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:47:41.744747Z",
     "start_time": "2023-02-08T02:47:41.691617Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[-9.8225653e-01,  1.9490494e+00],\n",
       "       [-3.1628523e+00, -6.4753151e-01],\n",
       "       [ 3.6174815e+00, -2.1167760e+00],\n",
       "       [-5.5663157e+00, -6.2317848e-03],\n",
       "       [-1.7511344e+00, -3.6932726e+00],\n",
       "       [ 2.4414740e+00, -9.9153006e-01],\n",
       "       [-4.6595273e+00, -1.0493579e+01],\n",
       "       [-3.0498884e+00,  1.7762201e+00],\n",
       "       [-6.9430470e-03,  1.6245110e+00],\n",
       "       [ 1.3217969e+00, -1.9404722e+00]], dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vmapped_affine = jax.vmap(affine)\n",
    "vmapped_affine(xs, ws, bs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It worked! \n",
    "\n",
    "\n",
    "## What if I want to vmap on one axis?\n",
    "\n",
    "Now if you might be yelling at me for having separate w and b for each x, especially if you're someone who is used to affine transformation (actually linear but all linear transformations are affine anyway) from machine learning or pytorch's nn.Linear. Fret not. We can have a single w and b for all xs. Actually let's write a simple linear discriminator (or classification) model, that takes a vector / batch of vector as inputs and gives a dim 2 vector as output as class probabilities. \n",
    "\n",
    "(Certainly not going to train this!, just the inference part. Besides I don't have any labels :P )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:47:44.063519Z",
     "start_time": "2023-02-08T02:47:44.061289Z"
    }
   },
   "outputs": [],
   "source": [
    "# using the same w, b values from before\n",
    "params = (w, b)\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def model(params, x):\n",
    "    w, b = params\n",
    "    logits = x @ w.T + b\n",
    "    \n",
    "    return jax.nn.softmax(logits, axis=-1) # convert to probability distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay but we have to keep only one value of params for all x, right? You can tell vmap to ignore params. How? using `in_axes`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:47:47.585579Z",
     "start_time": "2023-02-08T02:47:46.868653Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0.0507002 , 0.94929975],\n",
       "       [0.02904104, 0.970959  ],\n",
       "       [0.03516598, 0.964834  ],\n",
       "       [0.150932  , 0.84906805],\n",
       "       [0.09174331, 0.9082567 ],\n",
       "       [0.2507983 , 0.74920166],\n",
       "       [0.88580555, 0.11419442],\n",
       "       [0.9345587 , 0.06544128],\n",
       "       [0.95485044, 0.0451496 ],\n",
       "       [0.04433183, 0.9556682 ]], dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vmapped_model = jax.vmap(model, in_axes=(None, 0))\n",
    "vmapped_model(params, xs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`in_axes` lets you mention which params of the function you would like to be \"vmapped\". If you wanted to pass as `w, b, x` instead of `params, x`; you could've written\n",
    "\n",
    "```python\n",
    "vmapped_model = jax.vmap(model, in_axes=(None, None, 0))\n",
    "```\n",
    "\n",
    "A value of None will tell vmap to ignore that particular parameter. (Maintain order, you can randomly assign None !). For everything else, you can mention on which dimension should the array be vectorised. I've used 0 here. You can use something else for more complex tasks. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.11"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}