Julia_Jax_TF.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Julia vs. Jax vs. TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import jit, grad, jacfwd, jacrev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing Julia interpreter. This may take some time...\n"
     ]
    }
   ],
   "source": [
    "%load_ext julia.magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%julia\n",
    "using Zygote, ForwardDiff, BenchmarkTools, LinearAlgebra"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit\n",
    "def f(x):\n",
    "    for i in range(1,100):\n",
    "        x = x / (x + 1)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 95.1 ms, sys: 8.67 ms, total: 104 ms\n",
      "Wall time: 117 ms\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/marius/miniconda3/lib/python3.7/site-packages/jax/lib/xla_bridge.py:116: UserWarning: No GPU/TPU found, falling back to CPU.\n",
      "  warnings.warn('No GPU/TPU found, falling back to CPU.')\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.01, dtype=float32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# First evaluation where JAX is compiling the code\n",
    "%time f(1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "138 µs ± 6.06 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit f(2.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Same thing for the gradients we need to compile them\n",
    "g = jit(grad(f))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 540 ms, sys: 14 ms, total: 554 ms\n",
      "Wall time: 544 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.0001, dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time g(1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "149 µs ± 4.33 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit g(2.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Julia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%julia\n",
    "function f(x)\n",
    "    for i in 1:100\n",
    "        x = x / (x + 1)\n",
    "    end\n",
    "    return x\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  332.207 ns (0 allocations: 0 bytes)\n"
     ]
    }
   ],
   "source": [
    "%%julia\n",
    "@btime f(1.);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  54.992 μs (1564 allocations: 51.27 KiB)\n"
     ]
    }
   ],
   "source": [
    "%%julia\n",
    "@btime Zygote.gradient(f, 1.);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  800.330 ns (4 allocations: 320 bytes)\n"
     ]
    }
   ],
   "source": [
    "%%julia\n",
    "@btime ForwardDiff.gradient(x -> f(x[1]), [1.]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TensorFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is surely not that fair since you're mostly measuring overhead, but at least it gives you a sense of the overhead everytime you go from Python <-> TF, which is non-existant in Julia, and probably much less in a `@jit`'ed Jax function. A TF expert should probably improve this code. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def f(x):\n",
    "    for i in range(1,100):\n",
    "        x = x / (x + 1)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def gradf(x):\n",
    "    with tf.GradientTape() as t:\n",
    "        t.watch(x)\n",
    "        out = f(x)\n",
    "    return t.gradient(out, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.convert_to_tensor(1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 802 ms, sys: 188 ms, total: 990 ms\n",
      "Wall time: 1.44 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float32, numpy=0.010000002>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# First evaluation might be slow\n",
    "%time f(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.15 ms ± 51.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit f(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.3 s, sys: 15.5 ms, total: 1.32 s\n",
      "Wall time: 1.33 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float32, numpy=9.9999976e-05>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Same thing for gradients\n",
    "%time gradf(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.56 ms ± 62.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit gradf(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Jamie's HZPT model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this model, we are computing its Jacobian matrix with shape (50, 3)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analytic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "50.9 µs ± 426 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "r = np.logspace(-1,2)\n",
    "\n",
    "def hzpt_g_analytic(r,A0,R,R1h):\n",
    "    A_grad = - np.exp(-r/R) / (4*np.pi*r*R**2) *(1 - (R/R1h)**2 * np.exp(-(R+R1h)*r/(R*R1h)))\n",
    "    R_grad = A0*(-(r- 2*R)/R**2 + 2*np.exp(-r*(1/R + 1/R1h))*r /R1h**2) * np.exp(-r/R) / (4*np.pi*r*R**2)\n",
    "    R1h_grad = A0* np.exp(-r*(1/R + 1/R1h))* R**2 *(r-2*R1h)/R1h**4 * np.exp(-r/R) / (4*np.pi*r*R**2)\n",
    "    return np.array([A_grad,R_grad,R1h_grad])\n",
    "\n",
    "%timeit hzpt_g_analytic(r, 750.,26.,2.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/marius/miniconda3/lib/python3.7/site-packages/jax/lax/lax.py:5190: UserWarning: Explicitly requested dtype float64 requested in asarray is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.\n",
      "  warnings.warn(msg.format(dtype, fun_name , truncated_dtype))\n",
      "/home/marius/miniconda3/lib/python3.7/site-packages/jax/lax/lax.py:5190: UserWarning: Explicitly requested dtype float64 requested in array is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.\n",
      "  warnings.warn(msg.format(dtype, fun_name , truncated_dtype))\n"
     ]
    }
   ],
   "source": [
    "r = jnp.logspace(-1,2)\n",
    "\n",
    "@jit\n",
    "def hzpt_f_jax(r,A0,R,R1h):\n",
    "    F2 = jnp.exp(-r/R) / (4*jnp.pi*r*R**2)\n",
    "    return -A0 * F2 * (1 - (R/R1h)**2 * jnp.exp(-(R+R1h)*r/(R*R1h)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "280 µs ± 12.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "@jit\n",
    "def hzpt_jacfwd_jax(r, A0, R, R1h):\n",
    "    return jacfwd(hzpt_f_jax, argnums=(1, 2, 3))(r, A0, R, R1h)\n",
    "\n",
    "hzpt_jacfwd_jax(r, 750., 26., 2.)\n",
    "%timeit hzpt_jacfwd_jax(r, 750., 26., 2.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50,)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50,)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hzpt_jacfwd_jax(r, 750., 26., 2.)[1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "260 µs ± 13.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "@jit\n",
    "def hzpt_jacrev_jax(r, A0, R, R1h):\n",
    "    return jacrev(hzpt_f_jax, argnums=(1, 2, 3))(r, A0, R, R1h)\n",
    "\n",
    "hzpt_jacrev_jax(r, 750., 26., 2.)\n",
    "%timeit hzpt_jacrev_jax(r, 750., 26., 2.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Julia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%julia\n",
    "const r = 10 .^ range(-1, 2, length=50);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%julia\n",
    "function hzpt_f_jax(r, A0, R, R1h)\n",
    "    F2 = exp(-r/R) / (4π*r*R^2)\n",
    "    return -A0 * F2 * (1 - (R/R1h)^2 * exp(-(R+R1h)*r/(R*R1h)))\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  1.136 μs (1 allocation: 496 bytes)\n"
     ]
    }
   ],
   "source": [
    "%%julia\n",
    "@btime hzpt_f_jax.(r, 750., 26., 2.);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  2.911 μs (5 allocations: 3.47 KiB)\n"
     ]
    }
   ],
   "source": [
    "%%julia\n",
    "@btime ForwardDiff.jacobian(((A0, R, R1h),) -> hzpt_f_jax.(r, A0, R, R1h), [750., 26., 2.]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50, 3)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%julia\n",
    "jac = ForwardDiff.jacobian(((A0, R, R1h),) -> hzpt_f_jax.(r, A0, R, R1h), [750., 26., 2.])\n",
    "size(jac)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also try reverse-mode, where note I added the `norm` since we need the gradient of a scalar function. With only 3 arguments, the overhead of reverse-mode makes this significantly worse than forward-mode, although still faster than either of the Jax gradients. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  75.158 μs (1122 allocations: 68.70 KiB)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.3614060423837588, 0.26535983236863103, -261.8240445770739)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%julia\n",
    "@btime Zygote.gradient((A0, R, R1h) -> norm(hzpt_f_jax.(r, A0, R, R1h)), 750., 26., 2.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's also the hessian of a scalar function using all forward-mode:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  11.068 μs (10 allocations: 7.91 KiB)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 1.62630326e-19,  3.53813110e-04, -3.49098726e-01],\n",
       "       [ 3.53813110e-04, -2.54414437e-02, -1.20104043e-01],\n",
       "       [-3.49098726e-01, -1.20104043e-01,  3.77066747e+02]])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%julia\n",
    "@btime ForwardDiff.hessian(((A0, R, R1h),) -> norm(hzpt_f_jax.(r, A0, R, R1h)), [750., 26., 2.])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (with Julia)",
   "language": "python",
   "name": "python3-jl"
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}