jax_intro.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ff678c14",
   "metadata": {},
   "source": [
    "# An Introduction to JAX"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2efb7a5a",
   "metadata": {},
   "source": [
    "# GPU\n",
    "\n",
    "This lecture was built using a machine with JAX installed and access to a GPU.\n",
    "\n",
    "To run this lecture on [Google Colab](https://colab.research.google.com/), click on the “play” icon top right, select Colab, and set the runtime environment to include a GPU.\n",
    "\n",
    "To run this lecture on your own machine, you need to install [Google JAX](https://github.com/google/jax).\n",
    "\n",
    "This lecture provides a short introduction to [Google JAX](https://github.com/google/jax).\n",
    "\n",
    "Let’s see if we have an active GPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56961882",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d2beff2",
   "metadata": {},
   "source": [
    "## JAX as a NumPy Replacement\n",
    "\n",
    "One way to use JAX is as a plug-in NumPy replacement. Let’s look at the\n",
    "similarities and differences."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2345f5e4",
   "metadata": {},
   "source": [
    "### Similarities\n",
    "\n",
    "The following import is standard, replacing `import numpy as np`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a6a93e8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58c293d6",
   "metadata": {},
   "source": [
    "Now we can use `jnp` in place of `np` for the usual array operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79f05efb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = jnp.asarray((1.0, 3.2, -1.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6df0d5e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1547a5a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(jnp.sum(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ece7d3d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(jnp.mean(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f48bf55",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(jnp.dot(a, a))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "195c5576",
   "metadata": {},
   "source": [
    "However, the array object `a` is not a NumPy array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbf75e7b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f87bade3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "type(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba23ced0",
   "metadata": {},
   "source": [
    "Even scalar-valued maps on arrays return JAX arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9001305e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jnp.sum(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d9f444",
   "metadata": {},
   "source": [
    "JAX arrays are also called “device arrays,” where term “device” refers to a\n",
    "hardware accelerator (GPU or TPU).\n",
    "\n",
    "(In the terminology of GPUs, the “host” is the machine that launches GPU operations, while the “device” is the GPU itself.)\n",
    "\n",
    "Operations on higher dimensional arrays are also similar to NumPy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ee7137",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = jnp.ones((2, 2))\n",
    "B = jnp.identity(2)\n",
    "A @ B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61430ca7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from jax.numpy import linalg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a04d433a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "linalg.inv(B)   # Inverse of identity is identity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19342ae1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "linalg.eigh(B)  # Computes eigenvalues and eigenvectors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9cb7d18",
   "metadata": {},
   "source": [
    "### Differences\n",
    "\n",
    "One difference between NumPy and JAX is that JAX currently uses 32 bit floats by default.\n",
    "\n",
    "This is standard for GPU computing and can lead to significant speed gains with small loss of precision.\n",
    "\n",
    "However, for some calculations precision matters.  In these cases 64 bit floats can be enforced via the command"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccccaee4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "097f4dd9",
   "metadata": {},
   "source": [
    "Let’s check this works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2fa3ef6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jnp.ones(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "940f6918",
   "metadata": {},
   "source": [
    "As a NumPy replacement, a more significant difference is that arrays are treated as **immutable**.\n",
    "\n",
    "For example, with NumPy we can write"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "757e3cea",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "a = np.linspace(0, 1, 3)\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dffb40b",
   "metadata": {},
   "source": [
    "and then mutate the data in memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d815bcb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a[0] = 1\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f99cd51",
   "metadata": {},
   "source": [
    "In JAX this fails:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abf65207",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = jnp.linspace(0, 1, 3)\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04bac5b2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a[0] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5db96f69",
   "metadata": {},
   "source": [
    "In line with immutability, JAX does not support inplace operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35103bd6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.array((2, 1))\n",
    "a.sort()\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23ee71e4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = jnp.array((2, 1))\n",
    "a_new = a.sort()\n",
    "a, a_new"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b19d2ac",
   "metadata": {},
   "source": [
    "The designers of JAX chose to make arrays immutable because JAX uses a\n",
    "functional programming style.  More on this below.\n",
    "\n",
    "Note that, while mutation is discouraged, it is in fact possible with `at`, as in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dffcc007",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = jnp.linspace(0, 1, 3)\n",
    "id(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79b20da6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a54c41ee",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.at[0].set(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a04eccd5",
   "metadata": {},
   "source": [
    "We can check that the array is mutated by verifying its identity is unchanged:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb4c08bd",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "id(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "720c642f",
   "metadata": {},
   "source": [
    "## Random Numbers\n",
    "\n",
    "Random numbers are also a bit different in JAX, relative to NumPy.  Typically, in JAX, the state of the random number generator needs to be controlled explicitly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a735f062",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import jax.random as random"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e853eae1",
   "metadata": {},
   "source": [
    "First we produce a key, which seeds the random number generator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5539019",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "key = random.PRNGKey(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c183d8b9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "type(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e27fe4d9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(key)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b04eb875",
   "metadata": {},
   "source": [
    "Now we can use the key to generate some random numbers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9354441a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x = random.normal(key, (3, 3))\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0611206b",
   "metadata": {},
   "source": [
    "If we use the same key again, we initialize at the same seed, so the random numbers are the same:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "671a15f9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "random.normal(key, (3, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a36f68ad",
   "metadata": {},
   "source": [
    "To produce a (quasi-) independent draw, best practice is to “split” the existing key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9651b0a8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "key, subkey = random.split(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "995b69fb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "random.normal(key, (3, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0c7c4b7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "random.normal(subkey, (3, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aec9b827",
   "metadata": {},
   "source": [
    "The function below produces `k` (quasi-) independent random `n x n` matrices using this procedure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da3379b0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def gen_random_matrices(key, n, k):\n",
    "    matrices = []\n",
    "    for _ in range(k):\n",
    "        key, subkey = random.split(key)\n",
    "        matrices.append(random.uniform(subkey, (n, n)))\n",
    "    return matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da13d542",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "matrices = gen_random_matrices(key, 2, 2)\n",
    "for A in matrices:\n",
    "    print(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c57d0ab0",
   "metadata": {},
   "source": [
    "One point to remember is that JAX expects tuples to describe array shapes, even for flat arrays.  Hence, to get a one-dimensional array of normal random draws we use `(len, )` for the shape, as in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9da6d3a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "random.normal(key, (5, ))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21499e0b",
   "metadata": {},
   "source": [
    "## JIT compilation\n",
    "\n",
    "The JAX just-in-time (JIT) compiler accelerates logic within functions by fusing linear\n",
    "algebra operations into a single optimized kernel that the host can\n",
    "launch on the GPU / TPU (or CPU if no accelerator is detected)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96794be7",
   "metadata": {},
   "source": [
    "### A first example\n",
    "\n",
    "To see the JIT compiler in action, consider the following function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddf250fb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    a = 3*x + jnp.sin(x) + jnp.cos(x**2) - jnp.cos(2*x) - x**2 * 0.4 * x**1.5\n",
    "    return jnp.sum(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23d3885e",
   "metadata": {},
   "source": [
    "Let’s build an array to call the function on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3099c05",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = 50_000_000\n",
    "x = jnp.ones(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65b7be30",
   "metadata": {},
   "source": [
    "How long does the function take to execute?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb7cb364",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time f(x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea206ec6",
   "metadata": {},
   "source": [
    ">**Note**\n",
    ">\n",
    ">Here, in order to measure actual speed, we use the `block_until_ready()` method\n",
    "to hold the interpreter until the results of the computation are returned from\n",
    "the device. This is necessary because JAX uses asynchronous dispatch, which\n",
    "allows the Python interpreter to run ahead of GPU computations.\n",
    "\n",
    "The code doesn’t run as fast as we might hope, given that it’s running on a GPU.\n",
    "\n",
    "But if we run it a second time it becomes much faster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99d1d964",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time f(x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d26b3d73",
   "metadata": {},
   "source": [
    "This is because the built in functions like `jnp.cos` are JIT compiled and the\n",
    "first run includes compile time.\n",
    "\n",
    "Why would JAX want to JIT-compile built in functions like `jnp.cos` instead of\n",
    "just providing pre-compiled versions, like NumPy?\n",
    "\n",
    "The reason is that the JIT compiler can specialize on the *size* of the array\n",
    "being used, which is helpful for parallelization.\n",
    "\n",
    "For example, in running the code above, the JIT compiler produced a version of `jnp.cos` that is\n",
    "specialized to floating point arrays of size `n = 50_000_000`.\n",
    "\n",
    "We can check this by calling `f` with a new array of different size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "538b4b13",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "m = 50_000_001\n",
    "y = jnp.ones(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54315724",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time f(y).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50d46be2",
   "metadata": {},
   "source": [
    "Notice that the execution time increases, because now new versions of\n",
    "the built-ins like `jnp.cos` are being compiled, specialized to the new array\n",
    "size.\n",
    "\n",
    "If we run again, the code is dispatched to the correct compiled version and we\n",
    "get faster execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c30dfe60",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time f(y).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39107c67",
   "metadata": {},
   "source": [
    "The compiled versions for the previous array size are still available in memory\n",
    "too, and the following call is dispatched to the correct compiled code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dee5eac1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time f(x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3e2ca54",
   "metadata": {},
   "source": [
    "### Compiling the outer function\n",
    "\n",
    "We can do even better if we manually JIT-compile the outer function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b73509",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f_jit = jax.jit(f)   # target for JIT compilation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf11370c",
   "metadata": {},
   "source": [
    "Let’s run once to compile it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17892eac",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f_jit(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be71373f",
   "metadata": {},
   "source": [
    "And now let’s time it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aab63a7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time f_jit(x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21639ffc",
   "metadata": {},
   "source": [
    "Note the speed gain.\n",
    "\n",
    "This is because the array operations are fused and no intermediate arrays are created.\n",
    "\n",
    "Incidentally, a more common syntax when targetting a function for the JIT\n",
    "compiler is"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e7e7fd1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def f(x):\n",
    "    a = 3*x + jnp.sin(x) + jnp.cos(x**2) - jnp.cos(2*x) - x**2 * 0.4 * x**1.5\n",
    "    return jnp.sum(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff4e871b",
   "metadata": {},
   "source": [
    "## Functional Programming\n",
    "\n",
    "From JAX’s documentation:\n",
    "\n",
    "*When walking about the countryside of Italy, the people will not hesitate to tell you that JAX has “una anima di pura programmazione funzionale”.*\n",
    "\n",
    "In other words, JAX assumes a functional programming style.\n",
    "\n",
    "The major implication is that JAX functions should be pure.\n",
    "\n",
    "A pure function will always return the same result if invoked with the same inputs.\n",
    "\n",
    "In particular, a pure function has\n",
    "\n",
    "- no dependence on global variables and  \n",
    "- no side effects  \n",
    "\n",
    "\n",
    "JAX will not usually throw errors when compiling impure functions but execution becomes unpredictable.\n",
    "\n",
    "Here’s an illustration of this fact, using global variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2926d877",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = 1  # global\n",
    "\n",
    "@jax.jit\n",
    "def f(x):\n",
    "    return a + x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dc5ead1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x = jnp.ones(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fce3ae70",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6356d46",
   "metadata": {},
   "source": [
    "In the code above, the global value `a=1` is fused into the jitted function.\n",
    "\n",
    "Even if we change `a`, the output of `f` will not be affected — as long as the same compiled version is called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b33270e7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea183a3d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0cdbe05",
   "metadata": {},
   "source": [
    "Changing the dimension of the input triggers a fresh compilation of the function, at which time the change in the value of `a` takes effect:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d44316b8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x = jnp.ones(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c4317f2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38df871f",
   "metadata": {},
   "source": [
    "Moral of the story: write pure functions when using JAX!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70b02711",
   "metadata": {},
   "source": [
    "## Gradients\n",
    "\n",
    "JAX can use automatic differentiation to compute gradients.\n",
    "\n",
    "This can be extremely useful for optimization and solving nonlinear systems.\n",
    "\n",
    "We will see significant applications later in this lecture series.\n",
    "\n",
    "For now, here’s a very simple illustration involving the function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "239000f3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return (x**2) / 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89f88e89",
   "metadata": {},
   "source": [
    "Let’s take the derivative:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0363215e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f_prime = jax.grad(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0470e5a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f_prime(10.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afe60253",
   "metadata": {},
   "source": [
    "Let’s plot the function and derivative, noting that $ f'(x) = x $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "470a9d3b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "x_grid = jnp.linspace(-4, 4, 200)\n",
    "ax.plot(x_grid, f(x_grid), label=\"$f$\")\n",
    "ax.plot(x_grid, [f_prime(x) for x in x_grid], label=\"$f'$\")\n",
    "ax.legend(loc='upper center')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4376d9c9",
   "metadata": {},
   "source": [
    "We defer further exploration of automatic differentiation with JAX until [Adventures with Autodiff](https://jax.quantecon.org/autodiff.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc6caefd",
   "metadata": {},
   "source": [
    "## Writing vectorized code\n",
    "\n",
    "Writing fast JAX code requires shifting repetitive tasks from loops to array processing operations, so that the JAX compiler can easily understand the whole operation and generate more efficient machine code.\n",
    "\n",
    "This procedure is called **vectorization** or **array programming**, and will be familiar to anyone who has used NumPy or MATLAB.\n",
    "\n",
    "In most ways, vectorization is the same in JAX as it is in NumPy.\n",
    "\n",
    "But there are also some differences, which we highlight here.\n",
    "\n",
    "As a running example, consider the function\n",
    "\n",
    "$$\n",
    "f(x,y) = \\frac{\\cos(x^2 + y^2)}{1 + x^2 + y^2}\n",
    "$$\n",
    "\n",
    "Suppose that we want to evaluate this function on a square grid of $ x $ and $ y $ points and then plot it.\n",
    "\n",
    "To clarify, here is the slow `for` loop version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08a54be3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def f(x, y):\n",
    "    return jnp.cos(x**2 + y**2) / (1 + x**2 + y**2)\n",
    "\n",
    "n = 80\n",
    "x = jnp.linspace(-2, 2, n)\n",
    "y = x\n",
    "\n",
    "z_loops = np.empty((n, n))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b48095b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "for i in range(n):\n",
    "    for j in range(n):\n",
    "        z_loops[i, j] = f(x[i], y[j])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58ca3a21",
   "metadata": {},
   "source": [
    "Even for this very small grid, the run time is extremely slow.\n",
    "\n",
    "(Notice that we used a NumPy array for `z_loops` because we wanted to write to it.)\n",
    "\n",
    "OK, so how can we do the same operation in vectorized form?\n",
    "\n",
    "If you are new to vectorization, you might guess that we can simply write"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9584cd4b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z_bad = f(x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56228187",
   "metadata": {},
   "source": [
    "But this gives us the wrong result because JAX doesn’t understand the nested for loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "749101dd",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z_bad.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b77742a7",
   "metadata": {},
   "source": [
    "Here is what we actually wanted:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fca78288",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z_loops.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a4b0c48",
   "metadata": {},
   "source": [
    "To get the right shape and the correct nested for loop calculation, we can use a `meshgrid` operation designed for this purpose:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53eec135",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x_mesh, y_mesh = jnp.meshgrid(x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15c205df",
   "metadata": {},
   "source": [
    "Now we get what we want and the execution time is very fast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3220d6a7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "z_mesh = f(x_mesh, y_mesh).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8b5046d",
   "metadata": {},
   "source": [
    "Let’s run again to eliminate compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db66e617",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "z_mesh = f(x_mesh, y_mesh).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f75e89c4",
   "metadata": {},
   "source": [
    "Let’s confirm that we got the right answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bfb8383",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jnp.allclose(z_mesh, z_loops)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b3964a7",
   "metadata": {},
   "source": [
    "Now we can set up a serious grid and run the same calculation (on the larger grid) in a short amount of time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "528c5793",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = 6000\n",
    "x = jnp.linspace(-2, 2, n)\n",
    "y = x\n",
    "x_mesh, y_mesh = jnp.meshgrid(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "490a1acb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "z_mesh = f(x_mesh, y_mesh).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0522e70",
   "metadata": {},
   "source": [
    "Let’s run again to get rid of compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbd79ae0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "z_mesh = f(x_mesh, y_mesh).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "880c382e",
   "metadata": {},
   "source": [
    "But there is one problem here: the mesh grids use a lot of memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed73b56e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x_mesh.nbytes + y_mesh.nbytes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9354c668",
   "metadata": {},
   "source": [
    "By comparison, the flat array `x` is just"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00f92c04",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x.nbytes  # and y is just a pointer to x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "201ac9ab",
   "metadata": {},
   "source": [
    "This extra memory usage can be a big problem in actual research calculations.\n",
    "\n",
    "So let’s try a different approach using [jax.vmap](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html)\n",
    "\n",
    "First we vectorize `f` in `y`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e7268a0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f_vec_y = jax.vmap(f, in_axes=(None, 0))  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "077d584a",
   "metadata": {},
   "source": [
    "In the line above, `(None, 0)` indicates that we are vectorizing in the second argument, which is `y`.\n",
    "\n",
    "Next, we vectorize in the first argument, which is `x`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07501b8e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f_vec = jax.vmap(f_vec_y, in_axes=(0, None))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb636d2c",
   "metadata": {},
   "source": [
    "With this construction, we can now call the function $ f $ on flat (low memory) arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b2c4e2e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "z_vmap = f_vec(x, y).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ca8c8ef",
   "metadata": {},
   "source": [
    "We run it again to eliminate compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a833c00f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "z_vmap = f_vec(x, y).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00eb63fc",
   "metadata": {},
   "source": [
    "The execution time is essentially the same as the mesh operation but we are using much less memory.\n",
    "\n",
    "And we produce the correct answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12943164",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jnp.allclose(z_vmap, z_mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1006d896",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31f03b5e",
   "metadata": {},
   "source": [
    "## Exercise 2.1\n",
    "\n",
    "In the Exercise section of [a lecture on Numba and parallelization](https://python-programming.quantecon.org/parallelization.html), we used Monte Carlo to price a European call option.\n",
    "\n",
    "The code was accelerated by Numba-based multithreading.\n",
    "\n",
    "Try writing a version of this operation for JAX, using all the same\n",
    "parameters.\n",
    "\n",
    "If you are running your code on a GPU, you should be able to achieve\n",
    "significantly faster execution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25c7bde9",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 2.1](https://jax.quantecon.org/#jax_intro_ex2)\n",
    "\n",
    "Here is one solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d784af4f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "M = 10_000_000\n",
    "\n",
    "n, β, K = 20, 0.99, 100\n",
    "μ, ρ, ν, S0, h0 = 0.0001, 0.1, 0.001, 10, 0\n",
    "\n",
    "@jax.jit\n",
    "def compute_call_price_jax(β=β,\n",
    "                           μ=μ,\n",
    "                           S0=S0,\n",
    "                           h0=h0,\n",
    "                           K=K,\n",
    "                           n=n,\n",
    "                           ρ=ρ,\n",
    "                           ν=ν,\n",
    "                           M=M,\n",
    "                           key=jax.random.PRNGKey(1)):\n",
    "\n",
    "    s = jnp.full(M, np.log(S0))\n",
    "    h = jnp.full(M, h0)\n",
    "    for t in range(n):\n",
    "        key, subkey = jax.random.split(key)\n",
    "        Z = jax.random.normal(subkey, (2, M))\n",
    "        s = s + μ + jnp.exp(h) * Z[0, :]\n",
    "        h = ρ * h + ν * Z[1, :]\n",
    "    expectation = jnp.mean(jnp.maximum(jnp.exp(s) - K, 0))\n",
    "        \n",
    "    return β**n * expectation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72713918",
   "metadata": {},
   "source": [
    "Let’s run it once to compile it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7eee23a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time \n",
    "compute_call_price_jax().block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89c5ea7d",
   "metadata": {},
   "source": [
    "And now let’s time it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75328311",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time \n",
    "compute_call_price_jax().block_until_ready()"
   ]
  }
 ],
 "metadata": {
  "date": 1729509368.9189074,
  "filename": "jax_intro.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "An Introduction to JAX"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}