{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Diff On Vectors"
   ],
   "id": "site-title"
  },
  {
   "cell_type": "markdown",
   "id": "87d8a7f6",
   "metadata": {},
   "source": [
    "## Grad on Vectors and Matrices\n",
    "\n",
    "In real life machine learning projects, you're going to work more with vector and matrix data instead of simple numbers. In fact, the modern day GPUs have built in hardware parts to make matrix and vector computations lightning fast, hence the appeal to treat everything in terms of matrix computations. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ffb2308c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:09:23.493251Z",
     "start_time": "2023-02-08T01:09:23.259003Z"
    }
   },
   "outputs": [],
   "source": [
    "import jax as J\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eb05fa4",
   "metadata": {},
   "source": [
    "Let's start with a vector and find it's derivative. \n",
    "\n",
    "$h(x) =\\begin{vmatrix}\n",
    "9x^{2}ln(x)\n",
    "\\\\ \n",
    "10x\n",
    "\\end{vmatrix}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4607cd3c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:09:24.905200Z",
     "start_time": "2023-02-08T01:09:24.903171Z"
    }
   },
   "outputs": [],
   "source": [
    "def h(x):\n",
    "    return jnp.array(\n",
    "        [\n",
    "            [9 * x ** 2 * jnp.log(x)],\n",
    "            [10 * x]\n",
    "        ]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "944f22c6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:09:26.134494Z",
     "start_time": "2023-02-08T01:09:25.897229Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 0.],\n",
       "       [10.]], dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h(1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b46ca52c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:09:27.431542Z",
     "start_time": "2023-02-08T01:09:27.427779Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 1)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h(1.).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b51b84ab",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:09:29.200927Z",
     "start_time": "2023-02-08T01:09:28.860848Z"
    }
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "Gradient only defined for scalar-output functions. Output had shape: (2, 1).",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mTypeError\u001b[39m                                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[6]\u001b[39m\u001b[32m, line 2\u001b[39m\n\u001b[32m      1\u001b[39m d_h = J.grad(h)\n\u001b[32m----> \u001b[39m\u001b[32m2\u001b[39m \u001b[43md_h\u001b[49m\u001b[43m(\u001b[49m\u001b[32;43m1.\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "    \u001b[31m[... skipping hidden 4 frame]\u001b[39m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/api.py:523\u001b[39m, in \u001b[36m_check_scalar\u001b[39m\u001b[34m(x)\u001b[39m\n\u001b[32m    521\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(aval, ShapedArray):\n\u001b[32m    522\u001b[39m   \u001b[38;5;28;01mif\u001b[39;00m aval.shape != ():\n\u001b[32m--> \u001b[39m\u001b[32m523\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(msg(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mhad shape: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00maval.shape\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m))\n\u001b[32m    524\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m    525\u001b[39m   \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(msg(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mhad abstract value \u001b[39m\u001b[38;5;132;01m{\u001b[39;00maval\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m))\n",
      "\u001b[31mTypeError\u001b[39m: Gradient only defined for scalar-output functions. Output had shape: (2, 1)."
     ]
    }
   ],
   "source": [
    "d_h = J.grad(h)\n",
    "d_h(1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52f97c0d",
   "metadata": {},
   "source": [
    "Why isn't this working? It's because grad works on scalars (non vector / non list / non matrix) values or  just plain numbers. So we can't run  grad here. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85e7e9d3",
   "metadata": {},
   "source": [
    "## Jacobian\n",
    "\n",
    "An easy workaround is to get the Jacobian of h.\n",
    "\n",
    "An Jacobian of a vector or matrix is another vector or matrix containing all the first order derivatives of that vector or matrix. Sounds complicated? You can find more about Jacobian Matrix [here](https://youtu.be/bohL918kXQk) .\n",
    "\n",
    "Note: Second order derivatives are also known as Hessian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b4a82036",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:09:34.349242Z",
     "start_time": "2023-02-08T01:09:33.908354Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[ 9.],\n",
       "       [10.]], dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jc = J.jacobian(h)\n",
    "jc(1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5db4fb9",
   "metadata": {},
   "source": [
    "Now let's verify this thing row by row. For the first row, \n",
    "\n",
    "$\\frac{d}{dx}(9x^{2}ln(x)) = 9x + 18xln(x)$\n",
    "\n",
    "Which for $x = 1$ becomes $9$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c86faed9",
   "metadata": {},
   "source": [
    "For the second row, \n",
    "\n",
    "$\\frac{d}{dx}(10x) = 10$\n",
    "\n",
    "Which will stay $10$ no matter what $x$ is."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c789c40",
   "metadata": {},
   "source": [
    "## Vector to Vector\n",
    "\n",
    "Now let's make this interesting. Let's take the derivative of a vector in terms of another vector. \n",
    "\n",
    "\n",
    "$$\n",
    "x = [x_1 x_2]\n",
    "$$\n",
    "\n",
    "$$\n",
    "g = \\begin{vmatrix}\n",
    "e^{-8x_{1} + 6} & \\\\ \n",
    "\\frac{1}{3}e^{10x_{2}} & \\\\ \n",
    "3x_{1}9x_{2}\n",
    "\\end{vmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5bbdd766",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:09:37.612924Z",
     "start_time": "2023-02-08T01:09:37.610813Z"
    }
   },
   "outputs": [],
   "source": [
    "def g(x):\n",
    "    return jnp.array(\n",
    "        [\n",
    "            # exp is e\n",
    "            # https://numpy.org/doc/stable/reference/generated/numpy.exp.html\n",
    "            [jnp.exp(-8 * x[0] + 6)],\n",
    "            [(1. / 3.) * jnp.exp(10 * x[1])],\n",
    "            [3 * x[0] * 9 * x[1]]\n",
    "        ]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c7b0e17e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:09:41.378890Z",
     "start_time": "2023-02-08T01:09:40.600871Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[[-1.0826823e+00,  0.0000000e+00]],\n",
       "\n",
       "       [[ 0.0000000e+00,  1.6172174e+09]],\n",
       "\n",
       "       [[ 5.4000000e+01,  2.7000000e+01]]], dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d_g = J.jacobian(g)\n",
    "\n",
    "x = jnp.array([1., 2.])\n",
    "d_g(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef737518",
   "metadata": {},
   "source": [
    "![doge](../images/doge.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42c74074",
   "metadata": {},
   "source": [
    "## Explaim Hooman!\n",
    "\n",
    "Okay, let's do an \"explaim\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d35c5a86",
   "metadata": {},
   "source": [
    "For the first row, again, \n",
    "\n",
    "$\\frac{d}{dx_{1}}(e^{-8x_{1} + 6}) = -8e^{-8x_{1} + 6}$\n",
    "\n",
    "$\\frac{d}{dx_{2}}(e^{-8x_{1} + 6}) = 0$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0164a708",
   "metadata": {},
   "source": [
    "Second row, \n",
    "\n",
    "$\\frac{d}{dx_{1}}(\\frac{1}{3}e^{10x_{2}}) = 0$\n",
    "\n",
    "$\\frac{d}{dx_{2}}(\\frac{1}{3}e^{10x_{2}}) = \\frac{10}{3}e^{10x_{2}}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a11d563",
   "metadata": {},
   "source": [
    "Third row\n",
    "\n",
    "$\\frac{d}{dx_{1}}(3x_{1}9x_{2}) = 27x_{2}$\n",
    "\n",
    "$\\frac{d}{dx_{2}}(3x_{1}9x_{2}) = 27x_{1}$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f1344ba",
   "metadata": {},
   "source": [
    "Since we have a vector with 2 elements, x1 and x2, the final result will look like this:\n",
    "\n",
    "$$\n",
    "\\begin{vmatrix}\n",
    "-8e^{-8x_{1} + 6} & 0 \\\\ \n",
    "0 & \\frac{10}{3}e^{10x_{2}} \\\\ \n",
    "27x_{2} &  27x_{1}\n",
    "\\end{vmatrix}\n",
    "$$\n",
    "\n",
    "The first column containes derivatives in terms of x1 and the second one from x2."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45a7e852",
   "metadata": {},
   "source": [
    "## Matrix?\n",
    "\n",
    "You can use jacobian() on matrices as well. The procedure is same. (I can't ensure about verifying derivatives of large matrices in a simple manner though!)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.11"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}