{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Jit"
   ],
   "id": "site-title"
  },
  {
   "cell_type": "markdown",
   "id": "c49b7230",
   "metadata": {},
   "source": [
    "### The magic sauce over numpy\n",
    "\n",
    "Transforms are what Jax adds on top of numpy: `jit`, `grad`, `pmap`, `vmap`.\n",
    "\n",
    "This notebook is about jit."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ce8614d",
   "metadata": {},
   "source": [
    "### jit()\n",
    "\n",
    "Jax calls itself accelerated numpy. Or a faster numpy. Or however you interprete it. jit is one of the key components in this acceleration. What jit does is that it keeps a copy of your functions in the cache which can be executed faster than a regular function call. But how? \n",
    "\n",
    "\n",
    "\n",
    "#### Translating computer programs (or a rather CS101 refresher)\n",
    "\n",
    "Computers only understand 1's and 0's (in other words, binary numbers). So when you try to run your code, it needs to be translated into binary numbers first. This process isn't straightforward. There are multiple levels of what CS junkies call abstractions. One possible abstraction hierarchy can be like this (or what it used to look like before [llvm](https://llvm.org/) and [clang](https://clang.llvm.org/) came). \n",
    "\n",
    "\n",
    "```\n",
    "-------> Your code \n",
    "        -------> **Translator**\n",
    "                -------> Assembler (based on the instruction Set that your CPU / GPU maker made, also known as ISA) \n",
    "                        -------> Binary (executable)\n",
    "```\n",
    "\n",
    "\n",
    "\n",
    "This translator part in the middle can be either a Compiler or an Interpreter. Compiler takes your entire code file and translates it at once. There's a whole field of study in CS on compiler design if you want to look into that abyss. Anyway, since compilers translate or compile the whole file at once, you can add various optimizations (memory, speed) and enforce checks at translation or compile time to make sure your code has reduced amount of errors. Compiled programs are also, by comparison, faster. (Video Games, Operating Systems are prime examples.) **AND you need to compile only once!** (terms and conditions apply)\n",
    "\n",
    "Python on the other hand is interpreted. Interpreter translates one line at a time and **you have to intereprete everytime** you want to run your code. The process is slow and you can't enforce the same checks and optimizations like a compiler. Everything is known when the code is encountered so errors can't be reduced beforehand. So when Jax tries to cache your code for faster recall, the regular python design holds it back. Instead, Jax uses a compiler called XLA, which compiles your python code and caches it.\n",
    "\n",
    "\n",
    "And oh, JIT means just in time compilation (thank me later :) )\n",
    "\n",
    "So let's check how much of a performance benefit jit adds on top of regular code!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "77c46ac5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:12:34.348010Z",
     "start_time": "2023-02-08T01:12:34.116136Z"
    }
   },
   "outputs": [],
   "source": [
    "import jax as J\n",
    "from jax import jit, random\n",
    "import jax.numpy as jnp\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "eb65fb7c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:12:35.714470Z",
     "start_time": "2023-02-08T01:12:35.305384Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([ 1.6359469 ,  0.8408094 ,  0.02212393, -0.52927685, -0.6866449 ,\n",
       "       -1.0878773 , -1.2598456 ,  1.6036302 , -0.9536853 ], dtype=float32)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "key = random.key(123)\n",
    "\n",
    "x = random.normal(key, shape=(9, ))\n",
    "x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "29a9bb95",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:12:42.052635Z",
     "start_time": "2023-02-08T01:12:37.961241Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "115 \u03bcs \u00b1 4.51 \u03bcs per loop (mean \u00b1 std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "def f(x):\n",
    "    return jnp.sin(x) + jnp.cos(x)\n",
    "\n",
    "%timeit f(x)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef095714",
   "metadata": {},
   "source": [
    "With jit?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e5a7a596",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:12:57.421263Z",
     "start_time": "2023-02-08T01:12:45.502012Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37.8 \u03bcs \u00b1 1.02 \u03bcs per loop (mean \u00b1 std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "f_with_jit = jit(f)\n",
    "\n",
    "%timeit f_with_jit(x)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7497cc1",
   "metadata": {},
   "source": [
    "See that wall time difference ? That's like real fast!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe078dbb",
   "metadata": {},
   "source": [
    "#### Alternate syntax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9cefd33f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:13:20.152293Z",
     "start_time": "2023-02-08T01:13:08.452267Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "35.1 \u03bcs \u00b1 269 ns per loop (mean \u00b1 std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "@jit\n",
    "def f(x):\n",
    "    return jnp.sin(x) + jnp.cos(x)\n",
    "\n",
    "%timeit f(x)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd1c5b11",
   "metadata": {},
   "source": [
    "### But it's not all sunshine....\n",
    "\n",
    "Well caching and speeding up is nice but if your code changes, jit will have to recompile it. This adds some overhead and repeated compilations can make code execution slow, for example inside a loop. Also, not everything in python is supported in jax transformations, such as jit. Your code should be side-effect free (for starters, not have pythonic `loops`, `print` statements and `if-else` conditionals) and jax enforces a very functional programming approach. As a result you may not be able to jit every part of your code. \n",
    "\n",
    "This is where jax is different from Tensorflow or Pytorch. You've to think which parts of your code are going to be covered by transformations, which parts only take care of computations and which parts take care of loops and conditions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "036864e3",
   "metadata": {},
   "source": [
    "#### What my python code becomes after XLA compilation?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2b12f4a7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:13:30.340764Z",
     "start_time": "2023-02-08T01:13:30.339058Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:f32[9]\u001b[39m. \u001b[34;1mlet\n",
       "    \u001b[39;22mb\u001b[35m:f32[9]\u001b[39m = jit[\n",
       "      name=f\n",
       "      jaxpr={ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:f32[9]\u001b[39m. \u001b[34;1mlet\n",
       "          \u001b[39;22mc\u001b[35m:f32[9]\u001b[39m = sin a\n",
       "          d\u001b[35m:f32[9]\u001b[39m = cos a\n",
       "          b\u001b[35m:f32[9]\u001b[39m = add c d\n",
       "        \u001b[34;1min \u001b[39;22m(b,) }\n",
       "    ] a\n",
       "  \u001b[34;1min \u001b[39;22m(b,) }"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from jax import make_jaxpr\n",
    "\n",
    "make_jaxpr(f)(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf05cc39",
   "metadata": {},
   "source": [
    "This is what XLA does to your code and what Jax will see while executing it from a cache when flagged with jit.\n",
    "\n",
    "#### Now what if you do include some side effects or code unwanted by jax, for example a print statement?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "453260f1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:21:22.940168Z",
     "start_time": "2023-02-08T01:21:22.935314Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JitTracer<float32[9]>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:f32[9]\u001b[39m b\u001b[35m:f32[9]\u001b[39m. \u001b[34;1mlet\n",
       "    \u001b[39;22mc\u001b[35m:f32[9]\u001b[39m = jit[\n",
       "      name=side_effect_func\n",
       "      jaxpr={ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:f32[9]\u001b[39m b\u001b[35m:f32[9]\u001b[39m. \u001b[34;1mlet\u001b[39;22m c\u001b[35m:f32[9]\u001b[39m = add a b \u001b[34;1min \u001b[39;22m(c,) }\n",
       "    ] a b\n",
       "  \u001b[34;1min \u001b[39;22m(c,) }"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jit\n",
    "def side_effect_func(x, y):\n",
    "    print(x)\n",
    "    x = x + y # yes this is also a side effect\n",
    "    return x\n",
    "\n",
    "make_jaxpr(side_effect_func)(x, x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8e8b7dc9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:21:33.181956Z",
     "start_time": "2023-02-08T01:21:33.177471Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:f32[9]\u001b[39m b\u001b[35m:f32[9]\u001b[39m. \u001b[34;1mlet\n",
       "    \u001b[39;22mc\u001b[35m:f32[9]\u001b[39m = jit[\n",
       "      name=without_side_effect\n",
       "      jaxpr={ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:f32[9]\u001b[39m b\u001b[35m:f32[9]\u001b[39m. \u001b[34;1mlet\u001b[39;22m c\u001b[35m:f32[9]\u001b[39m = add a b \u001b[34;1min \u001b[39;22m(c,) }\n",
       "    ] a b\n",
       "  \u001b[34;1min \u001b[39;22m(c,) }"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@jit\n",
    "def without_side_effect(x, y):\n",
    "    return x + y\n",
    "\n",
    "make_jaxpr(without_side_effect)(x, x)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c69ca1d",
   "metadata": {},
   "source": [
    "See that there's `JitTracer<float32[9]>` in the jaxpr output for the function with side effects. So instead of being a pure jax type, it now has a tracer type, which has to adapt to the induced side effects, the print statement and modification of x. \n",
    "\n",
    "This may not be alarm bells to you as long as you're getting correct results but this may result in inconsistent behavior acorss devices. (which jax wants to avoid at all costs). \n",
    "\n",
    "You can read more here about [jax insisting on pure functions](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions).\n",
    "\n",
    "TL;DR : Just know that Jax is functional at core and only likes pure functions in your code. Anything else will probably run but jax can't provide you any guarantee about correct or consistent results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What if I want some parameters of my function to be not compiled?\n",
    "\n",
    "You may encounter this scenario if your function has to build an array with a shape which is not known prior (supplied as a parameter). Consider the following function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([5, 7, 0, 7, 7, 7, 1, 1, 1, 8, 1, 7, 0, 2, 4, 7, 1, 7, 1, 1, 7, 6,\n",
       "       0, 0, 7, 6, 7, 7, 7, 0, 0, 0, 1, 7, 7, 7, 0, 7, 0, 7, 0, 8, 5, 3,\n",
       "       0, 7, 0, 7, 7, 7], dtype=int32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def multinomial_sample(rng_key, logits, n_samples):\n",
    "    return random.categorical(rng_key, logits, shape=[n_samples, ])\n",
    "\n",
    "key, subkey = random.split(key)\n",
    "multinomial_sample(subkey, x, 50)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "jit'ing this function is not allowed, since `n_samples` can change and the array returned from the function won't have a fixed size on each call. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "Shapes must be 1D sequences of concrete values of integer type, got [JitTracer<~int32[]>].\nIf using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.\nThe error occurred while tracing the function multinomial_sample at /tmp/ipykernel_2099963/2375378534.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument n_samples.",
     "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[13]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43mjit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmultinomial_sample\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\u001b[43msubkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m50\u001b[39;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[12]\u001b[39m\u001b[32m, line 2\u001b[39m, in \u001b[36mmultinomial_sample\u001b[39m\u001b[34m(rng_key, logits, n_samples)\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mmultinomial_sample\u001b[39m(rng_key, logits, n_samples):\n\u001b[32m----> \u001b[39m\u001b[32m2\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrandom\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcategorical\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrng_key\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlogits\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshape\u001b[49m\u001b[43m=\u001b[49m\u001b[43m[\u001b[49m\u001b[43mn_samples\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/random.py:1787\u001b[39m, in \u001b[36mcategorical\u001b[39m\u001b[34m(key, logits, axis, shape, replace, mode)\u001b[39m\n\u001b[32m   1785\u001b[39m   shape = batch_shape\n\u001b[32m   1786\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1787\u001b[39m   shape = \u001b[43mcore\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcanonicalize_shape\u001b[49m\u001b[43m(\u001b[49m\u001b[43mshape\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1788\u001b[39m   _check_shape(\u001b[33m\"\u001b[39m\u001b[33mcategorical\u001b[39m\u001b[33m\"\u001b[39m, shape, batch_shape)\n\u001b[32m   1789\u001b[39m shape_prefix = shape[:\u001b[38;5;28mlen\u001b[39m(shape)-\u001b[38;5;28mlen\u001b[39m(batch_shape)]\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/core.py:2017\u001b[39m, in \u001b[36mcanonicalize_shape\u001b[39m\u001b[34m(shape, context)\u001b[39m\n\u001b[32m   2015\u001b[39m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m:\n\u001b[32m   2016\u001b[39m   \u001b[38;5;28;01mpass\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m2017\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m _invalid_shape_error(shape, context)\n",
      "\u001b[31mTypeError\u001b[39m: Shapes must be 1D sequences of concrete values of integer type, got [JitTracer<~int32[]>].\nIf using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.\nThe error occurred while tracing the function multinomial_sample at /tmp/ipykernel_2099963/2375378534.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument n_samples."
     ]
    }
   ],
   "source": [
    "jit(multinomial_sample)(subkey, x, 50)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The solution here is to tell jit that you want it to ignore the n_samples parameter. For that you have to use `static_argnums`. \n",
    "\n",
    "You can read more at: https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([5, 7, 0, 7, 7, 7, 1, 1, 1, 8, 1, 7, 0, 2, 4, 7, 1, 7, 1, 1, 7, 6,\n",
       "       0, 0, 7, 6, 7, 7, 7, 0, 0, 0, 1, 7, 7, 7, 0, 7, 0, 7, 0, 8, 5, 3,\n",
       "       0, 7, 0, 7, 7, 7], dtype=int32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jit(multinomial_sample, static_argnums=2)(subkey, x, 50)"
   ]
  }
 ],
 "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": 5
}