docs/tutorial_math/differentiation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b55233d4",
   "metadata": {},
   "source": [
    "# Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbc1e2d7",
   "metadata": {},
   "source": [
    "In this section, we are going to talk about how to make auto differentiation on your functions and class objects with 'jax' backend. In nowadays machine learning systems, computing and using gradients are common in various situations. So, we try to understand \n",
    "\n",
    "- how to  calculate derivatives of arbitrary complex functions, \n",
    "- how to compute high-order gradients. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e0ae6076",
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math.jax as bm\n",
    "\n",
    "bp.math.use_backend('jax')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29a369e5",
   "metadata": {},
   "source": [
    "All autodiff functions in BrainPy support pure functions and class objects. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df6b31f4",
   "metadata": {},
   "source": [
    "## ``grad()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d289c868",
   "metadata": {},
   "source": [
    "``brainpy.math.jax.grad()`` takes a function/object and returns a new function which computes the gradient of the original function/object. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56075f51",
   "metadata": {},
   "source": [
    "### Pure functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b4cccc",
   "metadata": {},
   "source": [
    "For pure function, the gradient is taken with respect to the first argument: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "45352485",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(a, b):\n",
    "    return a * 2 + b\n",
    "\n",
    "grad_f1 = bm.grad(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6009405f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(2., dtype=float32)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grad_f1(2., 1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c06f4f4e",
   "metadata": {},
   "source": [
    "However, this can be controlled via the `argnums` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "58aa6fbc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray(2., dtype=float32), DeviceArray(1., dtype=float32))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grad_f2 = bm.grad(f, argnums=(0, 1))\n",
    "\n",
    "grad_f2(2., 1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f0874ef",
   "metadata": {},
   "source": [
    "### Class objects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00906f22",
   "metadata": {},
   "source": [
    "For a class object or a class bound function, the gradient is taken with respect to the provided ``grad_vars`` argument: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "acc95d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class F(bp.Base):\n",
    "    def __init__(self):\n",
    "        super(F, self).__init__()\n",
    "        self.a = bm.TrainVar(bm.ones(1))\n",
    "        self.b = bm.TrainVar(bm.ones(1))\n",
    "\n",
    "    def __call__(self, c):\n",
    "        ab = self.a * self.b\n",
    "        ab2 = ab * 2\n",
    "        vv = ab2 + c\n",
    "        return vv.mean()\n",
    "    \n",
    "f = F()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18d64bc3",
   "metadata": {},
   "source": [
    "The ``grad_vars`` can be a JaxArray, or a list/tuple/dict of JaxArray. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "30484eab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'F0.a': TrainVar(DeviceArray([2.], dtype=float32)),\n",
       " 'F0.b': TrainVar(DeviceArray([2.], dtype=float32))}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.grad(f, grad_vars=f.train_vars())(10.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fa99d3ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[TrainVar(DeviceArray([2.], dtype=float32)),\n",
       " TrainVar(DeviceArray([2.], dtype=float32))]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.grad(f, grad_vars=[f.a, f.b])(10.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0847c77f",
   "metadata": {},
   "source": [
    "If there are values dynamically changed in the gradient function, you can provide them in the ``dyn_vars`` argument. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f77b4c0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class F2(bp.Base):\n",
    "    def __init__(self):\n",
    "        super(F2, self).__init__()\n",
    "        self.a = bm.TrainVar(bm.ones(1))\n",
    "        self.b = bm.TrainVar(bm.ones(1))\n",
    "\n",
    "    def __call__(self, c):\n",
    "        ab = self.a * self.b\n",
    "        ab = ab * 2\n",
    "        self.a.value = ab\n",
    "        return (ab + c).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a0cf62b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrainVar(DeviceArray([2.], dtype=float32))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f2 = F2()\n",
    "bm.grad(f2, dyn_vars=[f2.a], grad_vars=f2.b)(10.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6998ec7c",
   "metadata": {},
   "source": [
    "Also, if you are interested with the gradient of the input value, please use ``argnums`` argument. For this situation, calling the gradient function will return ``(grads_of_grad_vars, *grads_of_args)``. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "42c0dca2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class F3(bp.Base):\n",
    "    def __init__(self):\n",
    "        super(F3, self).__init__()\n",
    "        self.a = bm.TrainVar(bm.ones(1))\n",
    "        self.b = bm.TrainVar(bm.ones(1))\n",
    "\n",
    "    def __call__(self, c, d):\n",
    "        ab = self.a * self.b\n",
    "        ab = ab * 2\n",
    "        return (ab + c * d).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3fe1c9ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grads_of_gv : [TrainVar(DeviceArray([2.], dtype=float32)), TrainVar(DeviceArray([2.], dtype=float32))]\n",
      "grads_of_args : 3.0\n"
     ]
    }
   ],
   "source": [
    "f3 = F3()\n",
    "grads_of_gv, grad_of_arg0 = bm.grad(f3, grad_vars=[f3.a, f3.b], argnums=0)(10., 3.)\n",
    "\n",
    "print(\"grads_of_gv :\", grads_of_gv)\n",
    "print(\"grad_of_arg0 :\", grad_of_arg0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ba55cac6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grads_of_gv : [TrainVar(DeviceArray([2.], dtype=float32)), TrainVar(DeviceArray([2.], dtype=float32))]\n",
      "grad_of_arg0 : 3.0\n",
      "grad_of_arg1 : 10.0\n"
     ]
    }
   ],
   "source": [
    "f3 = F3()\n",
    "grads_of_gv, grad_of_arg0, grad_of_arg1 = bm.grad(f3, grad_vars=[f3.a, f3.b], argnums=(0, 1))(10., 3.)\n",
    "\n",
    "print(\"grads_of_gv :\", grads_of_gv)\n",
    "print(\"grad_of_arg0 :\", grad_of_arg0)\n",
    "print(\"grad_of_arg1 :\", grad_of_arg1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06491457",
   "metadata": {},
   "source": [
    "Actually, we recommend you to provide any dynamically changed variables (no matter them are updated in the gradient function) in the ``dyn_vars`` argument. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73cedb4d",
   "metadata": {},
   "source": [
    "### Auxiliary data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1469a67f",
   "metadata": {},
   "source": [
    "Usually, we want to get the value of the loss, or, we want to return some intermediate variables during the gradient computation. For them situation, users can set ``has_aux=True`` to return auxiliary data, and set ``return_value=True`` to return loss value. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a34a7e5a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grad:  TrainVar(DeviceArray([2.], dtype=float32))\n",
      "loss:  12.0\n"
     ]
    }
   ],
   "source": [
    "# return loss\n",
    "\n",
    "grad, loss = bm.grad(f, grad_vars=f.a, return_value=True)(10.)\n",
    "\n",
    "print('grad: ', grad)\n",
    "print('loss: ', loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "4a1ad862",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grad:  TrainVar(DeviceArray([1.], dtype=float32))\n",
      "aux_data:  (JaxArray(DeviceArray([1.], dtype=float32)), JaxArray(DeviceArray([2.], dtype=float32)))\n"
     ]
    }
   ],
   "source": [
    "class F4(bp.Base):\n",
    "    def __init__(self):\n",
    "        super(F4, self).__init__()\n",
    "        self.a = bm.TrainVar(bm.ones(1))\n",
    "        self.b = bm.TrainVar(bm.ones(1))\n",
    "\n",
    "    def __call__(self, c):\n",
    "        ab = self.a * self.b\n",
    "        ab2 = ab * 2\n",
    "        loss = (ab + c).mean()\n",
    "        return loss, (ab, ab2)\n",
    "    \n",
    "\n",
    "f4 = F4()\n",
    "    \n",
    "# return intermediate values\n",
    "grad, aux_data = bm.grad(f4, grad_vars=f4.a, has_aux=True)(10.)\n",
    "\n",
    "print('grad: ', grad)\n",
    "print('aux_data: ', aux_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33d2c322",
   "metadata": {},
   "source": [
    "**Note**: Any function wants to compute gradients through ``brainpy.math.jax.grad()`` must return a scalar value. Otherwise an error will raise. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "ea6a89f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'TypeError'> Gradient only defined for scalar-output functions. Output was [0. 0.].\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    bm.grad(lambda x: x)(bm.zeros(2))\n",
    "except Exception as e:\n",
    "    print(type(e), e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d08e3753",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "JaxArray(DeviceArray([0.5, 0.5], dtype=float32))"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# this is right\n",
    "bm.grad(lambda x: x.mean())(bm.zeros(2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1542356e",
   "metadata": {},
   "source": [
    "If you want to take gradients for a vector-output values, please use ``brainpy.math.jax.jacobian()`` function. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ca257d2",
   "metadata": {},
   "source": [
    "## ``jacobian()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f68747a3",
   "metadata": {},
   "source": [
    "Coming soon."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3a79f1f",
   "metadata": {},
   "source": [
    "## ``hessian()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d7da5a8",
   "metadata": {},
   "source": [
    "Coming soon."
   ]
  }
 ],
 "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": "279.273px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}