06_vmap.ipynb

{
 "cells": [
  {
   "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 the typical python fashion, 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 yada yada) and make your computation way faster. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's consider a simple affine transformation over some arrays.If you don't know what an affine transformation is, try refereshing your memory where you have seen this equation:\n",
    "\n",
    "$$\n",
    "y = xW^{T} + b\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "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": 3,
   "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": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:46:45.522829Z",
     "start_time": "2023-02-08T02:46:45.479643Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([-1.2992545,  6.513557 ], dtype=float32)"
      ]
     },
     "execution_count": 4,
     "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": 5,
   "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": 6,
   "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[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [6], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \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[0;31m[... skipping hidden 11 frame]\u001b[0m\n",
      "Cell \u001b[0;32mIn [2], line 3\u001b[0m, in \u001b[0;36maffine\u001b[0;34m(x, w, b)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;129m@jax\u001b[39m\u001b[38;5;241m.\u001b[39mjit\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21maffine\u001b[39m(x, w, b):\n\u001b[0;32m----> 3\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m@\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mw\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mT\u001b[49m \u001b[38;5;241m+\u001b[39m b\n",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "File \u001b[0;32m~/anaconda3/envs/jaxenv/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:4934\u001b[0m, in \u001b[0;36m_defer_to_unrecognized_arg.<locals>.deferring_binary_op\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m   4932\u001b[0m args \u001b[38;5;241m=\u001b[39m (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[1;32m   4933\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(other, _accepted_binop_types):\n\u001b[0;32m-> 4934\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mbinary_op\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   4935\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(other, _rejected_binop_types):\n\u001b[1;32m   4936\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124munsupported operand type(s) for \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mopchar\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m: \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   4937\u001b[0m                   \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(args[\u001b[38;5;241m0\u001b[39m])\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m and \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(args[\u001b[38;5;241m1\u001b[39m])\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
      "    \u001b[0;31m[... skipping hidden 5 frame]\u001b[0m\n",
      "File \u001b[0;32m~/anaconda3/envs/jaxenv/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:2979\u001b[0m, in \u001b[0;36mmatmul\u001b[0;34m(a, b, precision)\u001b[0m\n\u001b[1;32m   2977\u001b[0m a \u001b[38;5;241m=\u001b[39m lax\u001b[38;5;241m.\u001b[39msqueeze(a, \u001b[38;5;28mtuple\u001b[39m(a_squeeze))\n\u001b[1;32m   2978\u001b[0m b \u001b[38;5;241m=\u001b[39m lax\u001b[38;5;241m.\u001b[39msqueeze(b, \u001b[38;5;28mtuple\u001b[39m(b_squeeze))\n\u001b[0;32m-> 2979\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[43mlax\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdot_general\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2980\u001b[0m \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[43mndim\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;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[43mndim\u001b[49m\u001b[43m(\u001b[49m\u001b[43mb\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m-\u001b[39;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[1;32m   2981\u001b[0m \u001b[43m  \u001b[49m\u001b[43mprecision\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprecision\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2982\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m lax\u001b[38;5;241m.\u001b[39mtranspose(out, perm)\n",
      "    \u001b[0;31m[... skipping hidden 7 frame]\u001b[0m\n",
      "File \u001b[0;32m~/anaconda3/envs/jaxenv/lib/python3.10/site-packages/jax/_src/lax/lax.py:2543\u001b[0m, in \u001b[0;36m_dot_general_shape_rule\u001b[0;34m(lhs, rhs, dimension_numbers, precision, preferred_element_type)\u001b[0m\n\u001b[1;32m   2540\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m core\u001b[38;5;241m.\u001b[39msymbolic_equal_shape(lhs_contracting_shape, rhs_contracting_shape):\n\u001b[1;32m   2541\u001b[0m   msg \u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdot_general requires contracting dimensions to have the same \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   2542\u001b[0m          \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mshape, got \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m and \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m-> 2543\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(msg\u001b[38;5;241m.\u001b[39mformat(lhs_contracting_shape, rhs_contracting_shape))\n\u001b[1;32m   2545\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _dot_general_shape_computation(lhs\u001b[38;5;241m.\u001b[39mshape, rhs\u001b[38;5;241m.\u001b[39mshape, dimension_numbers)\n",
      "\u001b[0;31mTypeError\u001b[0m: dot_general requires contracting dimensions to have the same shape, got (10,) and (2,)."
     ]
    }
   ],
   "source": [
    "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": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:47:41.744747Z",
     "start_time": "2023-02-08T02:47:41.691617Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 2.3045597 ,  0.98218215],\n",
       "             [ 1.6167651 , -2.1813567 ],\n",
       "             [-3.882108  , -1.0218511 ],\n",
       "             [-0.12499112, -0.6038313 ],\n",
       "             [-2.508927  ,  0.8527329 ],\n",
       "             [-0.86322856,  3.2404249 ],\n",
       "             [-0.5951458 ,  1.4468247 ],\n",
       "             [-0.06645101, -0.23431951],\n",
       "             [-0.41957128, -7.0053816 ],\n",
       "             [-1.878496  , -2.0947793 ]], dtype=float32)"
      ]
     },
     "execution_count": 7,
     "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": 8,
   "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": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T02:47:47.585579Z",
     "start_time": "2023-02-08T02:47:46.868653Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[6.6032819e-04, 9.9933964e-01],\n",
       "             [7.1236588e-02, 9.2876339e-01],\n",
       "             [2.5964549e-01, 7.4035454e-01],\n",
       "             [9.0411073e-01, 9.5889255e-02],\n",
       "             [4.9198368e-01, 5.0801635e-01],\n",
       "             [9.8948145e-01, 1.0518541e-02],\n",
       "             [8.3628774e-01, 1.6371234e-01],\n",
       "             [9.8934537e-01, 1.0654581e-02],\n",
       "             [9.8608810e-01, 1.3911908e-02],\n",
       "             [4.2095959e-02, 9.5790404e-01]], dtype=float32)"
      ]
     },
     "execution_count": 9,
     "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\". So say for example 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. "
   ]
  }
 ],
 "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.10.8"
  },
  "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
  },
  "vscode": {
   "interpreter": {
    "hash": "85f6f424c393e7c95d987d0a64d03300802113ad2045bd4cdb9a90d6d84115f4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}