explain-04.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "74960ff4-9def-4b1d-9f37-caac6378adf3",
   "metadata": {},
   "source": [
    "# So what is the derivative of a matrix?\n",
    "## simple (!) (BUT ANNOYING!)\n",
    "### Assume that a Tensor will also have a `grad` field.\n",
    "### `Tensor.grad` carries with it the partial gradient of the final value with respect to that Tensor\n",
    "### Remember that we ultimately want to calculate the new .grad of both the left hand matrix and the right hand matrix\n",
    "### i.e. we need to calculate the outgoing gradient w.r.t. each $LH[i,j]$, i.e each value in  the left hand matrix and also $RH[i,j]$ for each value in the right hand matrix.\n",
    "\n",
    "### Assume `A` is (3,4) and `B` is (4,2). \n",
    "### if we do `C = A @ B`, the result of the matrix multiply `C` is a (3,2) matrix\n",
    "\n",
    "        C        = A                @ B\n",
    "        (3,2)      (3,4)              (4,2)  )\n",
    "        c00 c01    a00 a01 a02 a03    b00 b01\n",
    "        c10 c11  = a10 a11 a12 a13  @ b10 b11\n",
    "        c20 c21    a20 a21 a22 a23    b20 b21\n",
    "                                      b30 b31\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1c946b8-d0b6-4647-aff3-8c27b1ac948e",
   "metadata": {},
   "source": [
    "### Here it is, written out.\n",
    "\n",
    "        c00 = A[0,:] @ B[:,0]                    c01 = A[0,:] @ B[:,1]      \n",
    "            = a00 a01 a02 a03 @ b00                    a00 a01 a02 a03 @ b01\n",
    "                                b10                                      b11\n",
    "                                b20                                      b21\n",
    "                                b30                                      b31\n",
    "\n",
    "            = a00| |a01| |a02| |a03                  = a00| |a01| |a02| |a03\n",
    "               * |+| * |+| * |+| *                      * |+| * |+| * |+| * \n",
    "              b00| |b10| |b20| |b30                    b01| |b11| |b21| |b31\n",
    "\n",
    "            = a00*b00+a01*b10+a02*b20+a03*b30      = a00*b01+a01*b11+a02*b21*a03*b31\n",
    "\n",
    "        c10 = A[1,:] @ B[:,0]                    c11 = A[1,:] @ B[:,1]       \n",
    "            = a10 a11 a12 a13 @ b00                    a10 a11 a12 a13 @ b01\n",
    "                                b10                                      b11\n",
    "                                b20                                      b21\n",
    "                                b30                                      b31\n",
    "\n",
    "            = a10| |a11| |a12| |a13                  = a10| |a11| |a12| |a13\n",
    "               * |+| * |+| * |+| *                      * |+| * |+| * |+| * \n",
    "              b00| |b10| |b20| |b30                    b01| |b11| |b21| |b31\n",
    "\n",
    "            = a10*b00+a11*b10+a12*b20+a13*b30      = a10*b01+a11*b11+a12*b21*a13*b31\n",
    "\n",
    "        c20 = A[2,:] @ B[:,0]                    c21 = A[2,:] @ B[:,1]      \n",
    "            = a20 a21 a22 a23 @ b00                    a20 a21 a22 a23 @ b01\n",
    "                                b10                                      b11\n",
    "                                b20                                      b21\n",
    "                                b30                                      b31\n",
    "\n",
    "            = a20| |a21| |a22| |a23                  = a20| |a21| |a22| |a23\n",
    "               * |+| * |+| * |+| *                      * |+| * |+| * |+| * \n",
    "              b00| |b10| |b20| |b30                    b01| |b11| |b21| |b31\n",
    "\n",
    "            = a20*b00+a21*b10+a22*b20+a23*b30       = a20*b01+a21*b11+a22*b21*a23*b31\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c26a5468-95f0-4548-8396-cfc5fafa1ef2",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### So now, assuming we have the partial gradient for C\n",
    " - What is the gradient for A and B? (remember we're working backwards!)\n",
    " - Let's take a look at the expression `c00 = a00*b00+a01*b10+a02*b20+a03*b30`\n",
    " - and `c01 = a00*b01+a01*b11+a02*b21*a03*b31`\n",
    " - Notice that, for example, `a00` ONLY occurs at those two positions.\n",
    " - I.e. partial derivative `a00` must involve the partial derivatives of `C`, but only from those positions `c00` and `c01`\n",
    "  \n",
    "### So how do we calculate the following?\n",
    "\n",
    "       d(C)/d(A) =??                                          d(C)/d(B) == ?\n",
    "\n",
    " -  We need to calculate individual gradients for each entry in d(A) as well as d(B)\n",
    " -  i.e. in order to calculate d(C)/d(A), calculate d(C)/d(a00) ... d(C)/d(a23)\n",
    " -  We do that by figuring out how to distribute C.grad amongst A.grad and B.grad"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3476b72-ab61-4f70-958e-9fd8e1dbf940",
   "metadata": {},
   "source": [
    "### By the time we reached having to find the gradients for A and B, we must already have `C`'s gradient!\n",
    "\n",
    " - so to calculate d(C)/d(a00), see where `a00` occurs.\n",
    " - Since this is a linear multiply, the same behavior we saw in calculating the grad of the mult applies here!\n",
    " - since `a00` goes to 1, we know that `b00` and `b01` must be involved in the calculation. \n",
    " - if `C` was the last value in our computation chain, this means that the d(C)/d(a00) is `b00+b01`\n",
    " - OR `1*b00 + 1*b01`.. So what is the equivalent of the `1`'s in this case?\n",
    " - IT MUST BE the partial gradient of `C` at the locations where the `b00` and `b01` contribute to the value of `C`!\n",
    " - Those positions are `[0,0]` and `[0,1]`\n",
    " - The key thing to realize is that d(C)/d(c00) is simply `dc00` (i.e. `C.grad[0,0]`)\n",
    " - Therefore we replace `a00` with the cell coordinates `C`, where `a00` occurs.\n",
    " - i.e. at `[0,0]`, `a00` gets replaced with `dc00` and at `[0,1]`, `a00` is replaced with `dc01`\n",
    " - So, `d(C)/d(a00) = dc00*b00+dc01*b01`\n",
    " - Thus, the partial derivatives of d(C)/d(A[i,j]) is B[?] DOT C.grad[?] \n",
    "\n",
    "       d(C)/d(a00) = dc00*b00+dc01*b01                        d(C)/d(b00) = a00*dc00+a10*dc10+a20*dc20                \n",
    "       d(C)/d(a01) = dc00*b10+dc01*b11                        d(C)/d(b01) = a10*dc00+a11*dc11+a20*dc21                \n",
    "       d(C)/d(a02) = dc00*b20+dc01*b21                        d(C)/d(b10) = a01*dc00+a11*dc10+a22*dc20                \n",
    "       d(C)/d(a03) = dc00*b20+dc01*b21                        d(C)/d(b11) = a01*dc01+a11*dc11+a21*dc21  = A.T @ d(C)  \n",
    "                                                              d(C)/d(b20) = a02*dc00+a12*dc10+a22*dc20                \n",
    "       d(C)/d(a10) = dc10*b00+dc11*b01                        d(C)/d(b21) = a02*dc01+a12*dc11+a22*dc21                \n",
    "       d(C)/d(a11) = dc10*b10+dc11*b11      = d(C) @ B.T      d(C)/d(b30) = a02*dc00+a13*dc10+a23*dc20                \n",
    "       d(C)/d(a12) = dc10*b20+dc11*b21                        d(C)/d(b31) = a03*dc01+a13*dc11!a23*dc21\t\t\n",
    "       d(C)/d(a13) = dc10*b30+dc11*b31\n",
    "\n",
    "       d(C)/d(a20) = dc20*b00+dc21*b01\n",
    "       d(C)/d(a21) = dc20*b10+dc21*b11\n",
    "       d(C)/d(a22) = dc20*b20+dc21*b21\n",
    "       d(C)/d(a23) = dc20*b30+dc21*b31\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04c8386e-3733-42ae-a0c9-90457b629b34",
   "metadata": {},
   "source": [
    "### Here is the original matrix multiply\n",
    "\n",
    "         C        = A                @ B\n",
    "         (3,2)      (3,4)              (4,2) \n",
    "         c00 c01    a00 a01 a02 a03    b00 b01\n",
    "         c10 c11  = a10 a11 a12 a13  @ b10 b11\n",
    "         c20 c21    a20 a21 a22 a23    b20 b21\n",
    "                                       b30 b31\n",
    "                                       \n",
    "#### we have d(C), which is the same shape as C\n",
    "##### and we want to calculate d(C)/d(A) and d(C)/d(B)\n",
    "\n",
    "        d(C)/d(A) =          d(C)  @ B.T\n",
    "        (3,4)                (3,2) @ (2, 4)\n",
    "        da00 da01 da02 da03  = dc00 dc01   b00 b10 b20 b30 \n",
    "        da10 da11 da12 da13  = dc10 dc11   b01 b11 b21 b31\n",
    "        da20 da21 da22 da23  = dc20 dc21   \n",
    "\n",
    "        d(C)/d(B) =            A.T          @ d(C)\n",
    "        (4,2)                  (4,3)        @ (3,2)\n",
    "        db00 db01   =         a00 a10 a20    dc00 dc01\n",
    "        db10 db11             a01 a11 a21    dc10 dc11\n",
    "        db20 db21             a02 a12 a22    dc20 dc21\n",
    "        db30 db31             a03 a13 a23\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9e0375f-d3fd-4fa5-94e5-ca3648e03d31",
   "metadata": {},
   "source": [
    "# Here's a cheat:\n",
    "## Given the shapes of C (3,2), A (3,4), and B(4,2), and given that the local gradients of the C,A, and B, must have the same corresponding shape,...\n",
    "### Only way to get a d(C)/d(A) (3,4) matrix from d(C) and B is to d(C) @ B.T \n",
    "### Only way to get a d(C)/d(B) (4,2) matrix from d(C) and A is to A.T @ d(C)\n",
    "### Remember: d(C) is simply `C.grad`!"
   ]
  }
 ],
 "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}