docs/quickstart/jit_compilation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# JIT Compilation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "@[Chaoming Wang](mailto:adaduo@outlook.com)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The core idea behind BrainPy is the Just-In-Time (JIT) compilation. JIT compilation enables your Python code to be compiled into machine code \"just-in-time\" for execution. Subsequently, such transformed code can run at native machine code speed!\n",
    "\n",
    "Excellent JIT compilers such as [JAX](https://github.com/google/jax) and [Numba](https://github.com/numba/numba) are provided in Python. However, they are designed to work only on pure Python functions. While, in computational neuroscience, most models have too many parameters and variables, it's hard to manage and control model logic by only using functions. On the contrary, object-oriented programming (OOP) based on ``class`` in Python will make your coding more readable, controlable, flexible and modular. Therefore, it is necessary to support JIT compilation on class objects for programming in brain modeling. \n",
    "\n",
    "Here, in BrainPy, we provide JIT compilation interface for class objects, built on the top of JAX and Numba. In this section, we will talk about this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JIT in Numba and JAX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Numba](https://github.com/numba/numba) is specialized to optimize your native NumPy codes, including NumPy arrays, loops and condition controls, etc. It is a cross-platform library which can run on Windows, Linux, macOS, etc. The most wonderful thing is that numba can just-in-time compile your native Python loops (``for`` or ``while`` syntaxs) and condition controls (``if ... else ...``). This means that it supports your intutive Python programming. \n",
    "\n",
    "However, Numba is a lightweight JIT compiler, and is just suitable for small network models. For large networks, the parallel performance is poor. Futhermore, numba doesn't support `one code runs on multiple devices`. Same code cannot run on GPU targets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[JAX](https://github.com/google/jax) is a rising-star JIT compiler in Python scientific computing. It uses [XLA](https://www.tensorflow.org/xla) to JIT compile and run your NumPy programs. Same code can be deployed onto CPUs, GPUs and TPUs. Moreover, JAX supports automatic differentiation, which means you can train models through back-propagation. JAX prefers large network models, and has excellent parallel performance. \n",
    "\n",
    "However, JAX has intrinsic overhead, and is not suitable to run small networks. Moreover, JAX only supports Linux and macOS platforms. Windows users must install JAX on [WSL](https://docs.microsoft.com/en-us/windows/wsl/about) or compile JAX from source. Further, the coding in JAX is not very intutive. For example, \n",
    "\n",
    "- Doesn't support in-place mutating updates of arrays, like ``x[i] += y``, instead you should use `x = jax.ops.index_update(x, i, y)`\n",
    "- Doesn't support JIT compilation of your native loops and conditions, like\n",
    "```python\n",
    "arr = np.zeros(5)\n",
    "for i in range(arr.shape[0]):\n",
    "    arr[i] += 2.\n",
    "    if i % 2 == 0:\n",
    "        arr[i] += 1.\n",
    "```\n",
    "instead you should use \n",
    "```python\n",
    "arr = np.zeros(5)\n",
    "def loop_body(i, acc_arr):\n",
    "    arr1 = ops.index_update(acc_arr, i, acc_arr[i] + 2.)\n",
    "    return jax.lax.cond(i % 2 == 0, \n",
    "                        arr1,\n",
    "                        lambda arr1: ops.index_update(arr1, i, arr1[i] + 1),\n",
    "                        arr1,\n",
    "                        lambda arr1: arr1)\n",
    "arr = jax.lax.fori_loop(0, arr.shape[0], loop_body, arr)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What's more, both frameworks have poor support on class objects. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JIT compilation in BrainPy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to obtain an *intutive*, *flexible* and *high-performance* framework for brain modeling, in [BrainPy](https://github.com/PKU-NIP-Lab/BrainPy), we want to combine the advantages of both compilers together, and try to overcome the gotchas of each framework as much as possible (although we have not finished it). \n",
    "\n",
    "Specifically, we provide [BrainPy math module](../apis/math.rst) for \n",
    "\n",
    "- flexible switch between NumPy (Numba) and JAX backends\n",
    "- unified numpy-like array operations \n",
    "- unified ``ndarray`` data structure which supports in-place update\n",
    "- unified ``random`` APIs\n",
    "- powerful ``jit()`` compilation which supports functions and class objects both"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Backend Switch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To switch different backend, you can use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'numpy'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# switch to NumPy backend\n",
    "bm.use_backend('numpy')\n",
    "\n",
    "bm.get_backend_name()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'jax'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# switch to JAX backend\n",
    "bm.use_backend('jax')\n",
    "\n",
    "bm.get_backend_name()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In BrainPy, \"numpy\" and \"jax\" backends are interchangeable. Both backends nearly have the same APIs, and same codes can run on both backends."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Math Operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The APIs in ``brainpy.math`` module in each backend is much similar to APIs in original ``numpy``. The detailed comparison please see the [Comparison Table](../apis/math/comparison.rst). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, the **array creation** functions,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([[0., 0., 0.],\n",
       "                      [0., 0., 0.],\n",
       "                      [0., 0., 0.],\n",
       "                      [0., 0., 0.],\n",
       "                      [0., 0., 0.],\n",
       "                      [0., 0., 0.],\n",
       "                      [0., 0., 0.],\n",
       "                      [0., 0., 0.],\n",
       "                      [0., 0., 0.],\n",
       "                      [0., 0., 0.]], dtype=float32))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.zeros((10, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.arange(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([[1, 2],\n",
       "                      [3, 4]], dtype=int32))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = bm.array([[1,2], [3,4]])\n",
    "\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The **array manipulation** functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(4, dtype=int32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.max(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([1, 1, 2, 2, 3, 3, 4, 4], dtype=int32))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.repeat(x, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([[1, 1, 2, 2],\n",
       "                      [3, 3, 4, 4]], dtype=int32))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.repeat(x, 2, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The **random numbers** generation functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([[0.77751863, 0.45367956, 0.09705102, 0.35475922, 0.6678729 ],\n",
       "                      [0.81603515, 0.83256996, 0.29231536, 0.8294617 , 0.41171193],\n",
       "                      [0.37354684, 0.8089644 , 0.4921714 , 0.93098676, 0.4895897 ]],            dtype=float32))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.random.random((3, 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([[ 0.24033605,  1.6801527 ,  1.6716577 ,  0.19926837,\n",
       "                       -2.3778987 ],\n",
       "                      [ 0.58184516,  1.3625289 , -2.3364332 ,  2.082281  ,\n",
       "                        0.23409791]], dtype=float32))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = bm.random.normal(loc=0.0, scale=2.0, size=(2, 5))\n",
    "\n",
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The **linear algebra** functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([[ 1.4040264,  4.4052105, -3.0012088,  4.3638306, -1.9097029],\n",
       "                      [ 3.0483887, 10.490574 , -4.3307595,  8.926929 , -6.1973042]],            dtype=float32))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.dot(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(JaxArray(DeviceArray([-0.37228107+0.j,  5.3722816 +0.j], dtype=complex64)),\n",
       " JaxArray(DeviceArray([[-0.8245648 +0.j, -0.41597357+0.j],\n",
       "                       [ 0.56576747+0.j, -0.9093767 +0.j]], dtype=complex64)))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.linalg.eig(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The **discrete fourier transform** functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([ 3.2584137e-07+3.1786513e-08j,  8.0000000e+00+4.8023384e-07j,\n",
       "                      -3.2584137e-07+3.1786513e-08j, -1.6858739e-07+3.1786506e-08j,\n",
       "                      -3.8941437e-07-2.0663207e-07j,  2.3841858e-07-1.9411573e-07j,\n",
       "                       3.8941437e-07-2.0663207e-07j,  1.6858739e-07+3.1786506e-08j],            dtype=complex64))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.fft.fft(bm.exp(2j * bm.pi * bm.arange(8) / 8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([ 1.+0.j,  0.+1.j, -1.+0.j,  0.-1.j], dtype=complex64))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.fft.ifft(bm.array([0, 4, 0, 0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The full list of API implementation please see the [Comparison Table](../apis/math/comparison.rst)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### JIT for Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To take advantage of the JIT compilation, users just need to wrap their customized *functions* or *objects* into [bm.jit()](../apis/math/generated/brainpy.math.numpy.jit.rst) to instruct BrainPy to transform your codes into machine codes. \n",
    "\n",
    "\n",
    "Take the **pure functions** as the example. Here we try to implement a function of Gaussian Error Linear Unit,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "bm.use_backend('jax')\n",
    "\n",
    "\n",
    "def gelu(x):\n",
    "  sqrt = bm.sqrt(2 / bm.pi)\n",
    "  cdf = 0.5 * (1.0 + bm.tanh(sqrt * (x + 0.044715 * (x ** 3))))\n",
    "  y = x * cdf\n",
    "  return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first try the running speed without JIT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "332 µs ± 8.43 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "# jax backend, without JIT\n",
    "\n",
    "x = bm.random.random(100000)\n",
    "%timeit gelu(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function after JIT can significantly speed up. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "66.6 µs ± 135 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "# jax backend, with JIT\n",
    "\n",
    "gelu_jit = bm.jit(gelu)\n",
    "%timeit gelu_jit(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After siwtching to the 'numpy' backend, the result also applies.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "bm.use_backend('numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.52 ms ± 12.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "# numpy backend, without JIT\n",
    "\n",
    "x = bm.random.random(100000)\n",
    "%timeit gelu(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.91 ms ± 55.5 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "# numpy backend, with JIT\n",
    "\n",
    "gelu_jit = bm.jit(gelu)\n",
    "%timeit gelu_jit(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### JIT for Objects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, in BrainPy, JIT compilation can be carried on the **class objects**. Specifically, any instance of [brainpy.Base](../apis/generated/brainpy.base.Base.rst) object can be just-in-time compiled into machine codes. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try a simple example, which trains a Logistic regression classifier. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LogisticRegression(bp.Base):\n",
    "    def __init__(self, dimension):\n",
    "        super(LogisticRegression, self).__init__()\n",
    "\n",
    "        # parameters    \n",
    "        self.dimension = dimension\n",
    "    \n",
    "        # variables\n",
    "        self.w = bm.Variable(2.0 * bm.ones(dimension) - 1.3)\n",
    "\n",
    "    def __call__(self, X, Y):\n",
    "        u = bm.dot(((1.0 / (1.0 + bm.exp(-Y * bm.dot(X, self.w))) - 1.0) * Y), X)\n",
    "        self.w[:] = self.w - u\n",
    "        \n",
    "num_dim, num_points = 10, 20000000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we want to train the model weights (``self.w``) again and again. So, we mark the weights as ``bm.Variable``. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def benckmark(model, points, labels, num_iter=30, name=''):\n",
    "    t0 = time.time()\n",
    "    for i in range(num_iter):\n",
    "        model(points, labels)\n",
    "\n",
    "    print(f'{name} used time {time.time() - t0} s')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try to benchmark the 'numpy' backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "bm.use_backend('numpy')\n",
    "\n",
    "points = bm.random.random((num_points, num_dim))\n",
    "labels = bm.random.random(num_points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logistic Regression (without jit) used time 18.891679763793945 s\n"
     ]
    }
   ],
   "source": [
    "# numpy backend, without JIT\n",
    "\n",
    "lr1 = LogisticRegression(num_dim)\n",
    "\n",
    "benckmark(lr1, points, labels, name='Logistic Regression (without jit)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logistic Regression (with jit) used time 12.473472833633423 s\n"
     ]
    }
   ],
   "source": [
    "# numpy backend, with JIT\n",
    "\n",
    "lr2 = LogisticRegression(num_dim)\n",
    "lr2 = bm.jit(lr2)\n",
    "\n",
    "benckmark(lr2, points, labels, name='Logistic Regression (with jit)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same results can be obtained in 'jax' backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "bm.use_backend('jax')\n",
    "\n",
    "points = bm.random.random((num_points, num_dim))\n",
    "labels = bm.random.random(num_points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logistic Regression (without jit) used time 6.245944976806641 s\n"
     ]
    }
   ],
   "source": [
    "# jax backend, without JIT\n",
    "\n",
    "lr1 = LogisticRegression(num_dim)\n",
    "\n",
    "benckmark(lr1, points, labels, name='Logistic Regression (without jit)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Logistic Regression (with jit) used time 4.174307346343994 s\n"
     ]
    }
   ],
   "source": [
    "# jax backend, with JIT\n",
    "\n",
    "lr2 = LogisticRegression(num_dim)\n",
    "lr2 = bm.jit(lr2)\n",
    "\n",
    "benckmark(lr2, points, labels, name='Logistic Regression (with jit)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the JIT function showes its speed advantage."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, what's worth noting here is that:\n",
    "\n",
    "1. The dynamically changed variable (weight ``w``) is marked as a [bm.Variable](../apis/math/generated/brainpy.math.numpy.Variable.rst) (in `__init__()` function). \n",
    "2. The variable ``w``  is in-place updated with ``[:]`` indexing (in `__call__()` function)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above two things are all constraints in BrainPy for the JIT compilation of class objects. Other operations and coding styles are the same with the normal object-oriented programming in Python. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Mechanism of JIT in NumPy backend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "bm.use_backend('numpy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, **why must we in-place update the dynamically changed variables?**\n",
    "\n",
    "- First of all, in the compilation phase, a ``self.`` accessed variable which is not an instance of ``bm.Variable`` will be compiled as a static constant. For example, ``self.a = 1.`` will be compiled as a constant ``1.``. If you try to change the value of ``self.a``, it will not work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0\n"
     ]
    }
   ],
   "source": [
    "class Demo1(bp.Base):\n",
    "    def __init__(self):\n",
    "        super(Demo1, self).__init__()\n",
    "        \n",
    "        self.a = 1.\n",
    "    \n",
    "    def update(self, b):\n",
    "        self.a = b\n",
    "        \n",
    "\n",
    "d1 = Demo1()\n",
    "bm.jit(d1.update)(2.)\n",
    "print(d1.a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Second, all the variables you want to change during the function call must be labeled as ``bm.Variable``. Then during the JIT compilation period, these variables will be recompiled as arguments of the jitted functions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The recompiled function:\n",
      "-------------------------\n",
      "\n",
      "def update(b, Demo20_a=None):\n",
      "    Demo20_a = b\n",
      "\n",
      "\n",
      "The namespace of the above function:\n",
      "{}\n",
      "\n",
      "The recompiled function:\n",
      "-------------------------\n",
      "def new_update(b):\n",
      "  update(b, \n",
      "\t\tDemo20_a=Demo20.a.value,)\n",
      "\n",
      "The namespace of the above function:\n",
      "{'Demo20': <__main__.Demo2 object at 0x7fa9b055e2e0>,\n",
      " 'update': CPUDispatcher(<function update at 0x7fa9d01ff160>)}\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<function new_update(b)>"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Demo2(bp.Base):\n",
    "    def __init__(self):\n",
    "        super(Demo2, self).__init__()\n",
    "        \n",
    "        self.a = bm.Variable(1.)\n",
    "    \n",
    "    def update(self, b):\n",
    "        self.a = b\n",
    "        \n",
    "\n",
    "bm.jit(Demo2().update, show_code=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The original ``Demo2.update`` function is recompiled as ``update()`` function, with the dynamical variable ``a`` compiled as an argument ``Demo20_a``. Then, during the functional call (in the ``new_update()`` function), ``Demo20.a.value`` is passed to ``Demo20_a`` for the jitted ``update()`` function. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Third, as you can notice in the above source code of the recompiled function, the recompiled variable ``Demo20_a`` does not return. This means once the function finished its running, the computed value will disappear. Therefore, the dynamically changed variables must be in-place updated to hold their updated values. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable(2.)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Demo3(bp.Base):\n",
    "    def __init__(self):\n",
    "        super(Demo3, self).__init__()\n",
    "        \n",
    "        self.a = bm.Variable(1.)\n",
    "    \n",
    "    def update(self, b):\n",
    "        self.a[...] = b\n",
    "        \n",
    "\n",
    "d3 = Demo3()\n",
    "bm.jit(d3.update)(2.)\n",
    "d3.a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above simple demonstrations illustrate the core mechanism of the JIT compilation in NumPy backend. [bm.jit()](../apis/math/generated/brainpy.math.numpy.jit.rst) in NumPy backend can recursively compile your class objects. So, please try your models, and run it under the JIT accelerations.   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, the mechanism of JIT compilation of JAX backend is quite different. We will detail this in the upcoming tutorials. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### In-place operators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next, it is important to answer: **what are in-place operators?**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v = bm.arange(10)\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Actually, in-place operators include the following operations:\n",
    "\n",
    "1. **Indexing and slicing**. Like (More details please refer to [Array Objects Indexing](https://numpy.org/doc/stable/reference/arrays.indexing.html))\n",
    "  - Index: ``v[i] = a``\n",
    "  - Slice: ``v[i:j] = b``\n",
    "  - Slice the specific values: ``v[[1, 3]] = c``\n",
    "  - Slice all values, ``v[:] = d``, ``v[...] = e``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v[0] = 1\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v[1: 2] = 1\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v[[1, 3]] = 2\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v[:] = 0\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v[...] = bm.arange(10)\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. **Augmented assignment**. All augmented assignment are in-place operations, which include \n",
    "  - ``+=`` (add)\n",
    "  - ``-=`` (subtract)\n",
    "  - ``/=`` (divide)\n",
    "  - ``*=`` (multiply)\n",
    "  - ``//=`` (floor divide)\n",
    "  - ``%=`` (modulo)\n",
    "  - ``**=`` (power)\n",
    "  - ``&=`` (and)\n",
    "  - ``|=`` (or)\n",
    "  - ``^=`` (xor)\n",
    "  - ``<<=`` (left shift)\n",
    "  - ``>>=`` (right shift)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v += 1\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v *= 2\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v |= bm.random.randint(0, 2, 10)\n",
    "\n",
    "id (v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v **= 2.\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "140366785129408"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v >>= 2\n",
    "\n",
    "id(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More advanced usage please see our forthcoming tutorials. "
   ]
  }
 ],
 "metadata": {
  "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.9.5"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "245.76px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}