after_lectures/3/jax_autodiff_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Automatic differentiation with JAX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Main features\n",
    "\n",
    "- Numpy wrapper\n",
    "- Auto-vectorization\n",
    "- Auto-parallelization (SPMD paradigm)\n",
    "- Auto-differentiation\n",
    "- XLA backend and JIT support"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## How to compute gradient of your objective?\n",
    "\n",
    "- Define it as a standard Python function\n",
    "- Call ```jax.grad``` and voila!\n",
    "- Do not forget to wrap these functions with ```jax.jit``` to speed up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "- By default, JAX exploits single-precision numbers ```float32```\n",
    "- You can enable double precision (```float64```) by hands.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(5,) (5,)\n",
      "0.5431455433042264\n",
      "0.5431455433042264\n",
      "[[-0.17401344  0.09929537 -0.43481767  0.179563    0.54544231]\n",
      " [ 0.12356473 -0.07050838  0.30875849 -0.1275054  -0.38731164]\n",
      " [-0.08535699  0.04870632 -0.21328657  0.08807916  0.26755011]\n",
      " [-0.3445655   0.19661561 -0.8609862   0.35555423  1.08003499]\n",
      " [-0.20590302  0.11749217 -0.51450206  0.21246959  0.64539969]]\n",
      "(5, 1) (5,)\n",
      "0.5431455433042264\n"
     ]
    }
   ],
   "source": [
    "from jax.config import config\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "n = 5\n",
    "x = jax.random.normal(jax.random.PRNGKey(0), (n,))\n",
    "y = jax.random.normal(jax.random.PRNGKey(10), (n,))\n",
    "print(x.shape, y.shape)\n",
    "print(x @ y)\n",
    "print(x.T @ y)\n",
    "print(jnp.outer(x, y))\n",
    "print(x[:, None].shape, y.shape)\n",
    "print((x[None, :] @ y)[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "@jax.jit # Just-in-time compilation\n",
    "def f(x, A, b):\n",
    "    res = A @ x - b\n",
    "    res = jax.ops.index_update(res, 0, 100)\n",
    "#     y = res[res > 1]\n",
    "#     res[0] = 100\n",
    "    return res @ res\n",
    "\n",
    "gradf = jax.grad(f, argnums=0, has_aux=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Random numbers in JAX \n",
    "\n",
    "- JAX focuses on the reproducibility of the runs\n",
    "- Analogue of random seed is **the necessary argument** of all functions that generate something random\n",
    "- More details and references on the design of ```random``` submodule are [here](https://github.com/google/jax/blob/master/design_notes/prng.md)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "n = 1000\n",
    "x = jax.random.normal(jax.random.PRNGKey(0), (n, ))\n",
    "A = jax.random.normal(jax.random.PRNGKey(0), (n, n))\n",
    "b = jax.random.normal(jax.random.PRNGKey(0), (n, ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Check correctness 1388.1018567160188\n",
      "Compare speed\n",
      "Analytical gradient\n",
      "Grad function\n",
      "3.68 ms ± 483 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
      "Jitted grad function\n",
      "1.37 ms ± 160 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "print(\"Check correctness\", jnp.linalg.norm(gradf(x, A, b) - 2 * A.T @ (A @ x - b)))\n",
    "# print(gradf(x, A, b))\n",
    "print(\"Compare speed\")\n",
    "print(\"Analytical gradient\")\n",
    "# %timeit 2 * A.T @ (A @ x - b)\n",
    "print(\"Grad function\")\n",
    "%timeit gradf(x, A, b).block_until_ready()\n",
    "jit_gradf = jax.jit(gradf)\n",
    "print(\"Jitted grad function\")\n",
    "%timeit jit_gradf(x, A, b).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Check correctness 0.0\n",
      "Time for hessian\n",
      "95.7 ms ± 4.68 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n",
      "Emulate hessian and check correctness 0.0\n",
      "Time of emulating hessian\n",
      "100 ms ± 8.79 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "hess_func = jax.jit(jax.hessian(f))\n",
    "print(\"Check correctness\", jnp.linalg.norm(2 * A.T @ A - hess_func(x, A, b)))\n",
    "print(\"Time for hessian\")\n",
    "%timeit hess_func(x, A, b).block_until_ready()\n",
    "print(\"Emulate hessian and check correctness\", \n",
    "      jnp.linalg.norm(jax.jit(hess_func)(x, A, b) - jax.jacfwd(jax.jacrev(f))(x, A, b)))\n",
    "print(\"Time of emulating hessian\")\n",
    "hess_umul_func = jax.jit(jax.jacfwd(jax.jacrev(f)))\n",
    "%timeit hess_umul_func(x, A, b).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Summary\n",
    "\n",
    "- JAX is a simple and extensible tool in the problem where autodiff is crucial\n",
    "- JIT is a key to fast Python code\n",
    "- Input/output dimensions are important\n",
    "- Hessian matvec is faster than explicit hessian matrix by vector product"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Python (pytorch)",
   "language": "python",
   "name": "pytorch"
  },
  "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}