Copy_of_17_JAX_basic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iwCi-NRq-5US"
   },
   "source": [
    "Step 15: JAX for high-performance GPU computing\n",
    "----\n",
    "***\n",
    "\n",
    "In the upcoming steps, we will explore JAX, a powerful library for Python, and demonstrate how it can be utilized to either modify existing code or develop new code optimized for efficient GPU computing. This step will specifically focus on introducing the fundamental functions of the JAX library.\n",
    "\n",
    "[JAX](https://github.com/google/jax) is NumPy on the CPU, GPU, and TPU, with great automatic differentiation for high-performance computational research. Here is a short video to introduce JAX: https://www.youtube.com/watch?v=SFKEQs_Hu2c&t=94s\n",
    "\n",
    "\n",
    "\n",
    "**JAX** = **J**IT + **A**utoGrad + **X**LA\n",
    "\n",
    "**JIT**: just-in-time (JIT) compilation\n",
    "\n",
    "**AutoGrad**: automatically differentiation\n",
    "\n",
    "**XLA**: Accelerated Linear Algebra\n",
    "\n",
    "To effectively test the JAX code on GPUs, it's advisable to utilize **[Google Colab](https://colab.research.google.com/)**, which offers access to cloud-based GPU resources. This approach is particularly beneficial if you don't have local GPU capabilities. However, if you do have access to local GPU resources, they would also serve well for this purpose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "Nu6NHoLA-t6H"
   },
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "from jax import grad, jit, vmap\n",
    "from jax import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ELFmWmV5AnzZ",
    "outputId": "01903050-48f5-4e25-c75f-b6cad7bddb6a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.zeros(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3zBD4bI0A2Fv",
    "outputId": "3271b287-e4e2-4534-9570-2699f53711a0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.zeros(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IME2j9b4DQB7"
   },
   "source": [
    "### Multiplying Matrices\n",
    "We’ll be generating random data in the following examples. One big difference between NumPy and JAX is how you generate random numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xacALNaGA6MT",
    "outputId": "ed6b1123-fc43-4c62-ba34-731ea98e143e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.3721109   0.26423115 -0.18252768 -0.7368197  -0.44030377 -0.1521442\n",
      " -0.67135346 -0.5908641   0.73168886  0.5673026 ]\n"
     ]
    }
   ],
   "source": [
    "key = random.PRNGKey(0)\n",
    "x = random.normal(key, (10,))\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZkzLL4GJR9Wn"
   },
   "source": [
    "Let’s dive right in and multiply two big matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yAwnnDgTRaP1",
    "outputId": "fb450168-0750-4379-8d5b-5055d55769c6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.35 ms ± 2.22 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "size = 3000\n",
    "x = random.normal(key, (size,size), dtype = jnp.float32)\n",
    "%timeit jnp.dot(x, x.T).block_until_ready()  # runs on the GPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BVZ23o3RYYNi"
   },
   "source": [
    "We added that block_until_ready because JAX uses asynchronous execution by default (see Asynchronous dispatch)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "aLt8uhMNSJyN",
    "outputId": "3c9ed114-8f9b-4e9a-cd49-850549888de1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.13 s ± 669 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "y = np.random.rand(size, size)\n",
    "%timeit np.dot(y, y.T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n2_VwzPHYac_"
   },
   "source": [
    "You can see JAX numpy is much faster than numpy on computing matrix multiplication!\n",
    "\n",
    "JAX NumPy functions work on regular NumPy arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "A2ZKO908UUBB",
    "outputId": "12571da0-3e3c-4e70-e737-0cac3208e0db"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "36.7 ms ± 375 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "from jax._src.api import block_until_ready\n",
    "x = np.random.normal(size = (size,size)).astype(np.float32)\n",
    "%timeit jnp.dot(x, x.T).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8sBQU0qWaXrz"
   },
   "source": [
    "That’s slower because it has to transfer data to the GPU every time. You can ensure that an NDArray is backed by device memory using device_put()."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1qAwgtAeaV0e",
    "outputId": "3a102a74-3c9c-4d58-f51b-d10ec43d0ed1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.6 ms ± 177 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "from jax import device_put\n",
    "\n",
    "x = np.random.normal(size = (size,size)).astype(np.float32)\n",
    "x = device_put(x)\n",
    "%timeit jnp.dot(x, x.T).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4xpymGcTbMkP"
   },
   "source": [
    "The output of device_put() still acts like an NDArray, but it only copies values back to the CPU when they’re needed for printing, plotting, saving to disk, branching, etc. The behavior of device_put() is equivalent to the function jit(lambda x: x), but it’s faster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D_5hShjhbVfi"
   },
   "source": [
    "JAX is much more than just a GPU-backed NumPy. It also comes with a few program transformations that are useful when writing numerical code. For now, there are three main ones:\n",
    "\n",
    "jit(), for speeding up your code\n",
    "\n",
    "grad(), for taking derivatives\n",
    "\n",
    "vmap(), for automatic vectorization or batching.\n",
    "\n",
    "Let’s go over these, one-by-one. We’ll also end up composing these in interesting ways."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MAzd-X5ObiEp"
   },
   "source": [
    "### Using jit() to speed up functions\n",
    "JAX runs transparently on the GPU or TPU (falling back to CPU if you don’t have one). However, in the above example, JAX is dispatching kernels to the GPU one operation at a time. If we have a sequence of operations, we can use the @jit (just-in-time compile) decorator to compile multiple operations together using XLA (Accelerated Linear Algebra). Let’s try that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "m8UVijlga-EN",
    "outputId": "7d8f1e7f-ca19-4740-bc03-cc8f4f4816a2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "246 µs ± 8.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "@jit\n",
    "def selu(x, alpha = 1.67, lmbda = 1.05):\n",
    "  return lmbda * jnp.where(x>0, x, alpha * jnp.exp(x) - alpha)\n",
    "\n",
    "x = random.normal(key, (10000000,))\n",
    "%timeit selu(x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9V_U4n9YdOFF"
   },
   "source": [
    "The selu() function is called 10 million times. In this senario (a function is called many many times), we can speed it up with @jit, which will jit-compile the first time selu is called and will be cached thereafter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "LjizIwAMcZEZ",
    "outputId": "9b316281-8028-45bd-c564-cf775a039141"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "258 µs ± 4.11 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "selu_jit = jit(selu)\n",
    "%timeit selu_jit(x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qnJ_GGUMeN7A"
   },
   "source": [
    "Another way to use jit is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "oPu_WOr0d_R_",
    "outputId": "9e8b2299-f45e-433b-b742-2e352eac89a4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "258 µs ± 6.57 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "@jit\n",
    "def selu(x, alpha = 1.67, lmbda = 1.05):\n",
    "  return lmbda * jnp.where(x>0, x, alpha * jnp.exp(x) - alpha)\n",
    "\n",
    "x = random.normal(key, (10000000,))\n",
    "%timeit selu(x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PdINXy1wekqy"
   },
   "source": [
    "### Taking derivatives with grad()\n",
    "In addition to evaluating numerical functions, we also want to transform them. One transformation is automatic differentiation. In JAX, just like in Autograd, you can compute gradients with the grad() function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eIM-1oGRrxOp"
   },
   "source": [
    "With this configuration,config.update(\"jax_enable_x64\", True), JAX will now use float64 precision by default. Please note that the config.update(\"jax_enable_x64\", True) command should be called at the beginning of your script before importing JAX or modules that use JAX.\n",
    "\n",
    "Also, remember that using double precision comes with increased memory usage and computational cost, especially on GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "I623ALxWeLy3",
    "outputId": "09919644-4e03-4203-bbb3-9083a7ec248b"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-12-bce2302fe8d6>:1: DeprecationWarning: Accessing jax.config via the jax.config submodule is deprecated.\n",
      "  from jax.config import config\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.21 ms ± 391 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "[2.50000000e-01 1.96611933e-01 1.04993585e-01 4.51766597e-02\n",
      " 1.76627062e-02 6.64805667e-03 2.46650929e-03 9.10221180e-04\n",
      " 3.35237671e-04 1.23379350e-04 4.53958077e-05 1.67011429e-05\n",
      " 6.14413685e-06 2.26031919e-06 8.31527336e-07 3.05902133e-07\n",
      " 1.12535149e-07 4.13993738e-08 1.52299793e-08 5.60279637e-09\n",
      " 2.06115361e-09 7.58256042e-10 2.78946809e-10 1.02618796e-10\n",
      " 3.77513454e-11 1.38879439e-11 5.10908903e-12 1.87952882e-12\n",
      " 6.91440011e-13 2.54366565e-13 9.35762297e-14 3.44247711e-14\n",
      " 1.26641655e-14 4.65888615e-15 1.71390843e-15 6.30511676e-16\n",
      " 2.31952283e-16 8.53304763e-17 3.13913279e-17 1.15482242e-17\n",
      " 4.24835426e-18 1.56288219e-18 5.74952226e-19 2.11513104e-19\n",
      " 7.78113224e-20 2.86251858e-20 1.05306174e-20 3.87399763e-21\n",
      " 1.42516408e-21 5.24288566e-22 1.92874985e-22 7.09547416e-23\n",
      " 2.61027907e-23 9.60268005e-24 3.53262857e-24 1.29958143e-24\n",
      " 4.78089288e-25 1.75879220e-25 6.47023493e-26 2.38026641e-26\n",
      " 8.75651076e-27 3.22134029e-27 1.18506486e-27 4.35961000e-28\n",
      " 1.60381089e-28 5.90009054e-29 2.17052201e-29 7.98490425e-30\n",
      " 2.93748211e-30 1.08063928e-30 3.97544974e-31 1.46248623e-31\n",
      " 5.38018616e-32 1.97925988e-32 7.28129018e-33 2.67863696e-33\n",
      " 9.85415469e-34 3.62514092e-34 1.33361482e-34 4.90609473e-35\n",
      " 1.80485139e-35 6.63967720e-36 2.44260074e-36 8.98582594e-37\n",
      " 3.30570063e-37 1.21609930e-37 4.47377931e-38 1.64581143e-38\n",
      " 6.05460190e-39 2.22736356e-39 8.19401262e-40 3.01440879e-40\n",
      " 1.10893902e-40 4.07955867e-41 1.50078576e-41 5.52108228e-42\n",
      " 2.03109266e-42 7.47197234e-43 2.74878501e-43 1.01122149e-43]\n"
     ]
    }
   ],
   "source": [
    "from jax.config import config\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "from numpy.matrixlib import defmatrix\n",
    "\n",
    "def sum_logistic(x):\n",
    "  return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))\n",
    "\n",
    "x_small = jnp.arange(100.)\n",
    "y_small = jnp.linspace(0,3,100)\n",
    "\n",
    "derivative_fn = grad(sum_logistic)  # define a function which is the gradient of sum_logistic\n",
    "%timeit derivative_fn(x_small)\n",
    "print(derivative_fn(x_small))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QYIhu22PlRxC"
   },
   "source": [
    "Let’s verify with finite differences that our result is correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qAOMmTnqhXk_",
    "outputId": "75ed6955-7139-4829-892d-91a6c7b4d8d0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "208 ms ± 13.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      "[2.49999999e-01 1.96611936e-01 1.04993582e-01 4.51766624e-02\n",
      " 1.76627069e-02 6.64805810e-03 2.46650700e-03 9.10219455e-04\n",
      " 3.35241168e-04 1.23385746e-04 4.53965754e-05 1.66977543e-05\n",
      " 6.14619466e-06 2.25952590e-06 8.31335001e-07 3.05533376e-07\n",
      " 1.13686838e-07 4.26325641e-08 1.42108547e-08 7.10542736e-09\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00\n",
      " 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00]\n"
     ]
    }
   ],
   "source": [
    "def first_finite_differences(f, x):\n",
    "  eps = 1e-6\n",
    "  return jnp.array([(f(x + eps * v) - f(x - eps * v)) / (2 * eps)\n",
    "                   for v in jnp.eye(len(x))])\n",
    "\n",
    "%timeit first_finite_differences(sum_logistic, x_small)\n",
    "print(first_finite_differences(sum_logistic, x_small))\n",
    "# x = jnp.arange(10)\n",
    "# for v in jnp.eye(len(x)):\n",
    "#   print(v)\n",
    "# jnp.array([v for v in jnp.eye(len(x))])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gYFVh8HOsmaY"
   },
   "source": [
    "The results are almost the same, but you can see the bulit-in function grad() is much faster and has higher accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zweSWu7ptMAc"
   },
   "source": [
    "Taking derivatives is as easy as calling grad(). grad() and jit() compose and can be mixed arbitrarily. In the above example we jitted sum_logistic and then took its derivative. We can go further:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cb77dBnOlvNO",
    "outputId": "2102c513-f34d-460c-a8bc-6b4da52ac3b6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.03532558051623561\n"
     ]
    }
   ],
   "source": [
    "print(grad(jit(grad(jit(grad(sum_logistic)))))(1.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Im_lTVgctgGF"
   },
   "source": [
    "For more advanced autodiff, you can use jax.vjp() for reverse-mode vector-Jacobian products and jax.jvp() for forward-mode Jacobian-vector products. The two can be composed arbitrarily with one another, and with other JAX transformations. Here’s one way to compose them to make a function that efficiently computes full Hessian matrices:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "maMBKm8ztNYz"
   },
   "outputs": [],
   "source": [
    "from jax import jacfwd, jacrev\n",
    "def hessian(fun):\n",
    "  return jit(jacfwd(jacrev(fun)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nz2XRn0zti1k"
   },
   "source": [
    "### Auto-vectorization with vmap()\n",
    "\n",
    "JAX has one more transformation in its API that you might find useful: vmap(), the vectorizing map. It has the familiar semantics of mapping a function along array axes, but instead of keeping the loop on the outside, it pushes the loop down into a function’s primitive operations for better performance. When composed with jit(), it can be just as fast as adding the batch dimensions by hand.\n",
    "\n",
    "We’re going to work with a simple example, and promote matrix-vector products into matrix-matrix products using vmap(). Although this is easy to do by hand in this specific case, the same technique can apply to more complicated functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "T3Cu3quithe1",
    "outputId": "c2e0a316-c246-428f-f88d-dd73c3b331f2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 1 2 3]\n",
      "[[-0.53389115  0.84179134  0.81155729  0.05308707]\n",
      " [ 0.72478811 -0.53911566 -0.21932127  0.5509203 ]\n",
      " [ 0.16972549  1.19717228 -1.06094203  0.28213284]\n",
      " [-1.05431656  1.01875438 -0.42167228 -2.58898201]]\n",
      "[ 2.62416712  0.67500269 -0.07831328 -7.59153623]\n",
      "[[-0.          0.84179134  1.62311457  0.15926121]\n",
      " [ 0.         -0.53911566 -0.43864254  1.65276089]\n",
      " [ 0.          1.19717228 -2.12188407  0.84639851]\n",
      " [-0.          1.01875438 -0.84334457 -7.76694604]]\n",
      "[[0 0]\n",
      " [1 1]\n",
      " [2 2]\n",
      " [3 3]]\n",
      "[[ 2.62416712  2.62416712]\n",
      " [ 0.67500269  0.67500269]\n",
      " [-0.07831328 -0.07831328]\n",
      " [-7.59153623 -7.59153623]]\n",
      "[[ 2.62416712  2.62416712]\n",
      " [ 0.67500269  0.67500269]\n",
      " [-0.07831328 -0.07831328]\n",
      " [-7.59153623 -7.59153623]]\n"
     ]
    }
   ],
   "source": [
    "mat = random.normal(key, (150,100))\n",
    "batched_x = random.normal(key, (10,100))\n",
    "\n",
    "def apply_matrix(v):\n",
    "  return jnp.dot(mat, v)\n",
    "\n",
    "A = random.normal(key, (4,4))\n",
    "x = jnp.arange(4)\n",
    "y = jnp.arange(4)\n",
    "z = jnp.stack([x,y],1) # if axis = 0, jnp.dot(A,z) will have an error\n",
    "print(x)\n",
    "print(A)\n",
    "print(jnp.dot(A,x))  # jnp.dot is matrix multiplication. If x is a vector, it represents Ax\n",
    "print(A*x)           # A*x: A = [a1,a2,a3,a4] A*x = [a1*x1, a2*x2, a3*x3, a4*x4]\n",
    "\n",
    "print(z)\n",
    "print(jnp.dot(A,z))\n",
    "print(A @ z)         # A @ z is matrix multiplication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "EA63Sgt00kEu"
   },
   "outputs": [],
   "source": [
    "mat = random.normal(key, (150, 100))\n",
    "batched_x = random.normal(key, (10, 100))\n",
    "\n",
    "def apply_matrix(v):\n",
    "  return jnp.dot(mat, v)\n",
    "\n",
    "#apply_matrix(batched_x) #result in an error"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TgTY9p1pfyMm"
   },
   "source": [
    "Given a function such as apply_matrix, we can loop over a batch dimension in Python, but usually the performance of doing so is poor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xapnoi743A-3",
    "outputId": "7807befd-6f64-456c-9b73-19cf9b8ad931"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Naively batched\n",
      "4.1 ms ± 344 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "def naively_batched_apply_matrix(v_batched):\n",
    "  return jnp.stack([apply_matrix(v) for v in v_batched], 0)\n",
    "\n",
    "print('Naively batched')\n",
    "%timeit naively_batched_apply_matrix(batched_x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MpHCaEu5f8Q_"
   },
   "source": [
    "We know how to batch this operation manually. In this case, jnp.dot handles extra batch dimensions transparently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3Kc9jvUZf3_-",
    "outputId": "66545653-2eb4-4a57-9472-c74112871b50"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Manually batched\n",
      "131 µs ± 28.4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "@jit\n",
    "def batched_apply_matrix(v_batched):\n",
    "  return jnp.dot(v_batched, mat.T)\n",
    "\n",
    "print('Manually batched')\n",
    "%timeit batched_apply_matrix(batched_x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9JQ_FhGQf_pA"
   },
   "source": [
    "However, suppose we had a more complicated function without batching support. We can use vmap() to add batching support automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3G3QneyWf-Ev",
    "outputId": "2b3751cd-42fe-4f60-a14a-31bd66028a94"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Auto-vectorized with vmap\n",
      "147 µs ± 56.7 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "@jit\n",
    "def vmap_batched_apply_matrix(v_batched):\n",
    "  return vmap(apply_matrix)(v_batched)\n",
    "\n",
    "print('Auto-vectorized with vmap')\n",
    "%timeit vmap_batched_apply_matrix(batched_x).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MePQc8QRgEgA"
   },
   "source": [
    "Of course, vmap() can be arbitrarily composed with jit(), grad(), and any other JAX transformation.\n",
    "\n",
    "This is just a taste of what JAX can do. We’re really excited to see what you do with it!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xb70AbbuuUoA"
   },
   "source": [
    "### How to Think in JAX\n",
    "\n",
    "JAX provides a simple and powerful API for writing accelerated numerical code, but working effectively in JAX sometimes requires extra consideration. This document is meant to help build a ground-up understanding of how JAX operates, so that you can use it more effectively.\n",
    "\n",
    "### JAX vs. NumPy\n",
    "#### Key Concepts:\n",
    "\n",
    "JAX provides a NumPy-inspired interface for convenience.\n",
    "\n",
    "Through duck-typing, JAX arrays can often be used as drop-in replacements of NumPy arrays.\n",
    "\n",
    "Unlike NumPy arrays, JAX arrays are always immutable.\n",
    "\n",
    "NumPy provides a well-known, powerful API for working with numerical data. For convenience, JAX provides jax.numpy which closely mirrors the numpy API and provides easy entry into JAX. Almost anything that can be done with numpy can be done with jax.numpy.\n",
    "\n",
    "### NumPy, lax & XLA: JAX API layering\n",
    "#### Key Concepts:\n",
    "\n",
    "jax.numpy is a high-level wrapper that provides a familiar interface.\n",
    "\n",
    "jax.lax is a lower-level API that is stricter and often more powerful.\n",
    "\n",
    "All JAX operations are implemented in terms of operations in XLA – the Accelerated Linear Algebra compiler.\n",
    "\n",
    "If you look at the source of jax.numpy, you’ll see that all the operations are eventually expressed in terms of functions defined in jax.lax. You can think of jax.lax as a stricter, but often more powerful, API for working with multi-dimensional arrays.\n",
    "\n",
    "### To JIT or not to JIT\n",
    "\n",
    "#### Key Concepts:\n",
    "\n",
    "By default JAX executes operations one at a time, in sequence.\n",
    "\n",
    "Using a just-in-time (JIT) compilation decorator, sequences of operations can be optimized together and run at once.\n",
    "\n",
    "Not all JAX code can be JIT compiled, as it requires array shapes to be static & known at compile time.\n",
    "\n",
    "The fact that all JAX operations are expressed in terms of XLA allows JAX to use the XLA compiler to execute blocks of code very efficiently.\n",
    "\n",
    "**As a general rule of thumb, apply JIT (Just-In-Time) compilation only to functions that will be executed repeatedly in the program.**\n",
    "\n",
    "I encourage exploring further tutorials on JAX available on their [official website](https://jax.readthedocs.io/en/latest/) to deepen your understanding of its features and usage tips. However, for now, you have sufficient knowledge to proceed with converting our existing code to a JAX-enhanced version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7wUUihYZuWju"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "V100",
   "provenance": []
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}