explanations/comp_graph.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Backpropagation in depth\n",
    "\n",
    "In the [last lesson](https://github.com/VikParuchuri/zero_to_gpt/blob/master/explanations/rnn.ipynb), we learned how to create a recurrent neural network.  We now know how to build several network architectures using components like dense layers, softmax, and recurrent layers.\n",
    "\n",
    "We've been a bit loose with how we cover backpropagation, so that neural network architecture is easier to understand.  Backpropagation is how a neural network calculates how much to change each parameter in the network (the gradient).  Understanding how it works is important for tuning networks for performance, and writing fused kernels for GPUs.\n",
    "\n",
    "In this lesson, we'll do a deep dive into how backpropagation works.  We'll do this by building a computational graph to keep track of which changes we make to input data.\n",
    "\n",
    "A computational graph looks like this:\n",
    "\n",
    "![](images/comp_graph/comp_graph.png)\n",
    "\n",
    "It shows all the individual operations we performed (like multiplication) to modify the value of `X`, in order.  Keeping track of a computational graph is how we know how to reverse our operations to do backpropagation.\n",
    "\n",
    "This is the same way frameworks like PyTorch do forward and backward passes.  To create the computational graph, we'll make a miniature version of PyTorch!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Softmax function\n",
    "\n",
    "We'll first build a computational graph of the softmax function, then backpropagate through that graph to get the gradient against our inputs.  We introduced the softmax function in a previous lesson.  It's used to convert the output of a neural network into probabilities that can be used as predictions.  The softmax function is defined as:\n",
    "\n",
    "$$\\zeta=\\frac{e^{\\hat{y_{i}}}}{\\sum_{j=0}e^{\\hat{y_{j}}}}$$\n",
    "\n",
    "For each row of our neural network output, we raise $e$ to the power of our output value, then divide by the sum of $e$ raised to the power of each of the outputs for that row.\n",
    "\n",
    "The softmax function looks like this in code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def softmax_func(normalized):\n",
    "    raised = np.exp(normalized)\n",
    "    output = raised / np.sum(raised, axis=1).reshape(-1,1)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can test the softmax function using some fake data that we generate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5 rows and 3 columns of random numbers\n",
    "x = np.random.rand(5, 3)\n",
    "\n",
    "# Generate random correct labels for later\n",
    "# Exactly one label per row will be correct\n",
    "y = np.zeros_like(x)\n",
    "inds = (np.arange(0,y.shape[0]), np.random.randint(0, 3, size=y.shape[0]))\n",
    "y[inds] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`x` is our input to the softmax function.  It has `3` columns.  `y` is our target, where each row is a one-hot encoded vector.  The `1` will correspond to the correct label for each row.\n",
    "\n",
    "We can then apply the softmax function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.38110313, 0.34380373, 0.27509314],\n",
       "       [0.39553464, 0.41072091, 0.19374445],\n",
       "       [0.2744683 , 0.39579312, 0.32973858],\n",
       "       [0.37172812, 0.35272599, 0.27554589],\n",
       "       [0.46719982, 0.32717758, 0.20562259]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "normalized = x - np.max(x, axis=-1).reshape(-1,1)\n",
    "softmax_func(normalized)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above code, we subtract the maximum from each element in the row before passing the data into the softmax function.  This prevents numerical underflow or overflow.  Each [numeric type](https://numpy.org/doc/stable/user/basics.types.html) (float, integer, etc) can only hold a certain number of digits.  For example, floating point 16 can store 5 exponent bits, and ten digit bits (each bit is only base 2, so this is less than the same number of base-10 digits).  The maximum value we can store in `float16` is `65500`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "65500.0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check the maximum value we can assign to float16\n",
    "np.finfo('float16').max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/xz/9z84c__j28g8tg28bmcthjj00000gn/T/ipykernel_61885/2419725567.py:3: RuntimeWarning: overflow encountered in cast\n",
      "  a[0] = 6.55e5\n"
     ]
    }
   ],
   "source": [
    "# This is an example of numeric overflow, where we store more digits than float16 can hold\n",
    "a = np.array([0], dtype=np.float16)\n",
    "a[0] = 6.55e5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we raise $e$ to a very large or small number, we can generate a number that is too large to store in our specific data type.  Subtracting the max gives us the same end result, but reduces the risk of overflow.  Feel free to try the softmax out with and without subtracting the max to see how it works!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Staged Softmax\n",
    "\n",
    "Instead of computing the softmax derivative, we previously used the fact that the derivative of the softmax and negative log likelihood functions \"cancel out\", and end up with a derivative of $p-y$.  But what if we want to find the derivative ourselves?\n",
    "\n",
    "We can approach it analytically, and find the derivative of the entire function.  We can even use SymPy to help us do the derivation, like we did in an earlier lesson.  Another method is to break the softmax function apart into individual operations.  Each operation will make a single modification to the data:\n",
    "\n",
    "![](images/comp_graph/softmax_steps.svg)\n",
    "\n",
    "We perform 3 operations on the data:\n",
    "\n",
    "- Exp - we raise e to the power x.\n",
    "- Sum - we add up the $e^x$ values for each row.\n",
    "- Divide - we divide the $e^x$ values by the sums.\n",
    "\n",
    "Note that the output of `Exp` is passed to both the `Sum` and `Divide` operations.\n",
    "\n",
    "By breaking up the softmax this way, we can take the derivative of each individual piece instead of the whole function at once.  By the [chain rule](https://www.khanacademy.org/math/ap-calculus-ab/ab-differentiation-2-new/ab-3-1a/a/chain-rule-review), multiplying the derivative of each individual operation will result in the derivative of the whole function.  We used the chain rule in previous lessons to find the partial derivative of the loss with respect to the model weights and biases."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can build the forward pass of our staged softmax.  The derivative of multiplication is easier to calculate than division, so we'll swap some of our operations to remove the division.\n",
    "\n",
    "Luckily for us, raising a value `x` to the power `-1` is the same as taking `1/x`.  So instead of dividing `Exp/Sum`, we can do `Exp * Sum ^ -1`, leaving us with these operations:\n",
    "\n",
    "- Exp\n",
    "- Sum\n",
    "- Pow - we invert the sum by raising to the power `-1`\n",
    "- Multiply - we multiply the inverted sum and the exp values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "raised = np.exp(x) # step 1\n",
    "summed = np.sum(raised, axis=-1).reshape(-1,1) # step 2.  reshape so each row has 1 column.\n",
    "pow = summed ** -1 # step 3\n",
    "staged_softmax = raised * pow # step 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.38110313, 0.34380373, 0.27509314],\n",
       "       [0.39553464, 0.41072091, 0.19374445],\n",
       "       [0.2744683 , 0.39579312, 0.32973858],\n",
       "       [0.37172812, 0.35272599, 0.27554589],\n",
       "       [0.46719982, 0.32717758, 0.20562259]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "staged_softmax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our staged softmax has the exact same output as our original function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Softmax Derivative\n",
    "\n",
    "To get the derivative of the softmax, we need to reverse the operations we did before:\n",
    "\n",
    "![](images/comp_graph/softmax_steps_full_bwd.svg)\n",
    "\n",
    "To compute the full derivative, here are the steps we need to follow:\n",
    "\n",
    "1. Start at the final step.  Take in the loss gradient as the input, and multiply by the derivative of the final step.\n",
    "2. Pass the gradient to the previous operation.\n",
    "2. Continue calculating the derivative of each operation, and multiplying by the gradient.  Note that `Exp` is input to two operations, so it will sum both gradients.\n",
    "4. Continue until we reach the first operation.\n",
    "\n",
    "To calculate loss, we'll use negative log likelihood, which is $NLL = - \\sum_{i=0} y_{i} \\log p_{i}$.  Since $y$ is only non-zero at one position per row, this will only have a single nonzero term ($-y_{i} * \\log p_{i}$ where $i$ is the correct label where $y$ equals `1`).\n",
    "\n",
    "We can solve this derivative by breaking it into steps:\n",
    "\n",
    "* $\\log p_{i}$ - the derivative of a natural log (base e) is $\\frac{1}{p_{i}}$\n",
    "* $-y_{i} * \\log p_{i}$ - the derivative wrt p is $-y_{i}$\n",
    "\n",
    "So the derivative of NLL is $\\frac{1}{p} * -y_{i} = \\frac{-y_{i}}{p_{i}}$\n",
    "\n",
    "We'll use the negative log likelihood derivative $\\frac{\\partial L}{\\partial p}$ below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def nll_grad(y, pred):\n",
    "    return -1 * y / pred\n",
    "\n",
    "loss_grad = nll_grad(y, staged_softmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then calculate the softmax derivative by multiplying the derivatives of the individual operations.  Remember that a derivative is the rate of change of a function as we change the input.\n",
    "\n",
    "- Exp - the derivative of $e^x$ is $e^x$ (this is a very cool property of $e$!)\n",
    "- Sum - since a sum operation will combine input elements into one, we just distribute the gradient over all input elements.  A change to any of the input elements will have a direct impact on the output.\n",
    "- Pow - the derivative of $x^{-1}$ is $-1 * x^{-2}$.  More on this [here](https://www.khanacademy.org/math/old-ap-calculus-ab/ab-derivative-rules/ab-diff-negative-fraction-powers/a/power-rule-review).\n",
    "- Multiply - the derivatives of $x*y$ are $y$ wrt $x$ and $x$ wrt $y$.  This is because any change to $x$ is multiplied by $y$, and vice versa.  Thus the rate of change of $x$ is $y$, and vice versa.\n",
    "\n",
    "We can now create the backward pass of our staged softmax.  The backward pass will start with the loss gradient.  This will be a matrix showing how much we need to adjust each of the output values from our softmax to reduce our loss.  We can then compute gradients for each operation, ending with the gradient against the input, `x`.  If `x` was the output of a neural network, we would continue backpropagation at that point to adjust the network parameters.\n",
    "\n",
    "We'll name each gradient according to the step it is a gradient for, not the step it is coming from.  So `raised_grad` is the gradient on `raised`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Step 4 derivative\n",
    "raised_grad = loss_grad * pow\n",
    "pow_grad = loss_grad * raised\n",
    "pow_grad = np.sum(pow_grad, axis=-1).reshape(-1,1) # reshape gradient to match input data\n",
    "\n",
    "# Step 3\n",
    "summed_grad = (-1 * summed ** -2) * pow_grad\n",
    "\n",
    "# Step 2\n",
    "raised_grad_2 = np.ones_like(raised) * summed_grad # distribute gradient across inputs\n",
    "\n",
    "# Step 1\n",
    "raised_grad += raised_grad_2 # sum incoming gradients\n",
    "staged_softmax_grad = raised_grad * np.exp(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We did two things above that might be confusing.  The first is that we summed 2 gradients on raised.  This is because raised connects to 2 operations, and both have separate gradients.  Whenever this happens, we sum the gradients.\n",
    "\n",
    "The second is that we reshaped `pow_grad` to have a single column.  This is to match `pow`, which only had `1` column in the forward pass.  Whenever a gradient doesn't match the shape of the input data, we change the size of the gradient to match it.  This is because the gradient represents the partial derivative against the input data to the operation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare our result to the derivative of the softmax equation to make sure everything worked.  The derivative of the softmax is $S_{i}((i==j) - S_{j})$.  We take each element of a single row in the output of a softmax, like this:\n",
    "\n",
    "    [0.28, 0.25, 0.47]\n",
    "\n",
    "We then compare each element against each other element.  So we could start at element `0` (`.28`), and compare it to itself.  Then `i` is `0` and `j` is `0`.  So the equation is `.28 * (1 - .28)`.  When we then keep `i` the same, but move `j` to `1`.  The equation becomes `.28 * (0 - .25)`.  And so on, until we construct a matrix like this:\n",
    "\n",
    "![](images/comp_graph/softmax_deriv.svg)\n",
    "\n",
    "We then sum across the rows and multiply by the incoming gradient to get the partial derivative against the inputs.  We can define the softmax derivative in code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def softmax_grad_func(softmax, loss_grad):\n",
    "    output = np.zeros_like(softmax)\n",
    "    for i in range(softmax.shape[0]):\n",
    "        sm_row = softmax[i,:]\n",
    "        sm_grad = (-np.outer(sm_row, sm_row) + np.diag(sm_row.flatten()))\n",
    "        row_grad = sm_grad * loss_grad[i,:].reshape(1,-1)\n",
    "        output[i,:] = np.sum(row_grad, -1)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then compare our derivative by stages with the derivation.  The `np.allclose` function tells us if all the values in an array are close to another array.  We use this instead of `==` because there are small numerical differences in similar computations with numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "derived_softmax_grad = softmax_grad_func(staged_softmax, loss_grad)\n",
    "np.allclose(derived_softmax_grad, staged_softmax_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also compare the gradient we computed of the loss and softmax together with the derivative of both together, which is `p-y`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(staged_softmax_grad, staged_softmax - y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Operations\n",
    "\n",
    "Breaking the softmax into stages helped us understand the basic units of a computational graph.  But what if we don't want to have to type out all the code for the forward and backward pass every time?  It would be nice if we could only define the forward pass, and automatically have the backward pass happen.  This is how it works in frameworks like PyTorch.\n",
    "\n",
    "We can build our own version of PyTorch by individually defining each operation, then mixing and matching the operations to create a more complex equation.  By doing this, we can create a computational graph like this one:\n",
    "\n",
    "![](images/comp_graph/comp_graph.png)\n",
    "\n",
    "Each node in the graph will be a separate class that knows how to do a forward and backward pass.  So we can just execute the graph to run forward and backward passes.\n",
    "\n",
    "We can start out by defining the operations.  I've written a class called `Node`, which we can subclass to define each operation.  You can look at the code for `Node` if you want.  It gives us some nice methods for running the operations in a graph in order, both forward and backward:\n",
    "\n",
    "- `apply_fwd` - runs the forward pass up to the node it is called on.\n",
    "- `apply_bwd` - runs the backward pass from the node it is called on backwards.\n",
    "- `zero_grad` - zero out our gradient before running a backward pass.\n",
    "- `generate_graph` - helps us visualize the computational graph.\n",
    "- `generate_derivative_chains` - shows us the equation for calculating the partial derivative at a node.\n",
    "\n",
    "For each operation, we just implement the `forward` and `backward` methods, which take in input data, and pass them through the operation.  The `Node` class takes care of the rest.\n",
    "\n",
    "We'll use the same formulas for each operation that we used in our staged softmax earlier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import sys, os\n",
    "sys.path.append(os.path.abspath(\"../nnets\"))\n",
    "from graph import Node, Parameter, display_chain\n",
    "\n",
    "class Exp(Node):\n",
    "    def forward(self, x):\n",
    "        return np.exp(x) # raise e to the power x\n",
    "\n",
    "    def backward(self, grad):\n",
    "        x = self.cache[0] # Pull the x value used in the forward pass\n",
    "        return np.exp(x) * grad # multiply the incoming gradient by the derivative\n",
    "\n",
    "class Sum(Node):\n",
    "    def forward(self, x):\n",
    "        return np.sum(x, axis=-1).reshape(-1,1)\n",
    "\n",
    "    def backward(self, grad):\n",
    "        x = self.cache[0] # Pull the x value used in the forward pass\n",
    "        return np.ones_like(x) * grad # distribute the gradient over the input data shape\n",
    "\n",
    "class Pow(Node):\n",
    "    def forward(self, x, exponent):\n",
    "        return x ** exponent\n",
    "\n",
    "    def backward(self, grad):\n",
    "        x, exponent = self.cache # Pull the x and exponent values used in the forward pass\n",
    "        return grad * exponent * x ** (exponent - 1), 1\n",
    "\n",
    "class Multiply(Node):\n",
    "    def forward(self, x, y):\n",
    "        return x * y\n",
    "\n",
    "    def backward(self, grad):\n",
    "        x, y = self.cache # Pull the x and y values used in the forward pass\n",
    "        return grad * y, grad * x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see above, we've defined the 4 operations that we need for our softmax in code.\n",
    "\n",
    "We can now define our whole softmax operation as a computational graph.  When we initialize a `Node`, we pass in the nodes that feed into it.  If we're using the node to feed in data (like inputs), we use a special `Parameter` node:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# The input data for our network.  We set needs_grad=True so the gradient is calculated for this parameter in the backward pass.\n",
    "# Desc is a short description of what the data in this node is.\n",
    "X = Parameter(x, desc=\"X\", needs_grad=True)\n",
    "\n",
    "# Raise e to the power x.  Out is a description of the output of the node.\n",
    "raised = Exp(X, out=\"e^X\")\n",
    "# Sum the raised values.\n",
    "summed = Sum(raised, out=\"sum(e^X)\")\n",
    "\n",
    "# Define -1 as a parameter, so we can use it as an exponent.\n",
    "negative_one = Parameter(-1, desc=\"-1\", needs_grad=False)\n",
    "# Invert our sums\n",
    "inverted = Pow(summed, negative_one, out=\"1 / sum(e^X)\")\n",
    "# Multiply the inverted sums by e^X\n",
    "softmax = Multiply(raised, inverted, out=\"softmax(X)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now use the `generate_graph` method on the `softmax` node to visualize the computational graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 7.0.6 (20230106.0513)\n",
       " -->\n",
       "<!-- Title: fwd_pass Pages: 1 -->\n",
       "<svg width=\"614pt\" height=\"128pt\"\n",
       " viewBox=\"0.00 0.00 614.09 128.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 124)\">\n",
       "<title>fwd_pass</title>\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-124 610.09,-124 610.09,4 -4,4\"/>\n",
       "<!-- 4516507376 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>4516507376</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"563.84\" cy=\"-18\" rx=\"42.49\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"563.84\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">Multiply</text>\n",
       "</g>\n",
       "<!-- 4516759200 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>4516759200</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"127\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"127\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">Exp</text>\n",
       "</g>\n",
       "<!-- 4516759200&#45;&gt;4516507376 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>4516759200&#45;&gt;4516507376</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M154.22,-18C224.04,-18 415.52,-18 509.46,-18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"509.44,-21.5 519.44,-18 509.44,-14.5 509.44,-21.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"313.1\" y=\"-21.8\" font-family=\"Times,serif\" font-size=\"14.00\">e^X</text>\n",
       "</g>\n",
       "<!-- 4516758192 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>4516758192</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"240.3\" cy=\"-102\" rx=\"27.1\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"240.3\" y=\"-98.3\" font-family=\"Times,serif\" font-size=\"14.00\">Sum</text>\n",
       "</g>\n",
       "<!-- 4516759200&#45;&gt;4516758192 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>4516759200&#45;&gt;4516758192</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M145.39,-31.12C163.42,-44.73 191.82,-66.17 212.68,-81.91\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"210.36,-84.55 220.45,-87.78 214.58,-78.96 210.36,-84.55\"/>\n",
       "<text text-anchor=\"middle\" x=\"183.5\" y=\"-70.8\" font-family=\"Times,serif\" font-size=\"14.00\">e^X</text>\n",
       "</g>\n",
       "<!-- 4516758288 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>4516758288</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">X</text>\n",
       "</g>\n",
       "<!-- 4516758288&#45;&gt;4516759200 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>4516758288&#45;&gt;4516759200</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M54.26,-18C64.7,-18 76.9,-18 88.25,-18\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"88.15,-21.5 98.15,-18 88.15,-14.5 88.15,-21.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"77\" y=\"-21.8\" font-family=\"Times,serif\" font-size=\"14.00\">X</text>\n",
       "</g>\n",
       "<!-- 4516758528 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>4516758528</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"385.6\" cy=\"-48\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"385.6\" y=\"-44.3\" font-family=\"Times,serif\" font-size=\"14.00\">Pow</text>\n",
       "</g>\n",
       "<!-- 4516758528&#45;&gt;4516507376 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>4516758528&#45;&gt;4516507376</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M412.22,-43.64C438.44,-39.17 479.87,-32.12 512.64,-26.54\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"513.07,-30.02 522.34,-24.89 511.9,-23.12 513.07,-30.02\"/>\n",
       "<text text-anchor=\"middle\" x=\"467.1\" y=\"-42.8\" font-family=\"Times,serif\" font-size=\"14.00\">1 / sum(e^X)</text>\n",
       "</g>\n",
       "<!-- 4516758192&#45;&gt;4516758528 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>4516758192&#45;&gt;4516758528</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M265.14,-94.11C285.36,-87.3 315.06,-77 340.6,-67 344.46,-65.48 348.5,-63.84 352.49,-62.17\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"353.57,-65.51 361.39,-58.36 350.82,-59.07 353.57,-65.51\"/>\n",
       "<text text-anchor=\"middle\" x=\"313.1\" y=\"-89.8\" font-family=\"Times,serif\" font-size=\"14.00\">sum(e^X)</text>\n",
       "</g>\n",
       "<!-- 4516758240 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>4516758240</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"240.3\" cy=\"-48\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"240.3\" y=\"-44.3\" font-family=\"Times,serif\" font-size=\"14.00\">&#45;1</text>\n",
       "</g>\n",
       "<!-- 4516758240&#45;&gt;4516758528 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>4516758240&#45;&gt;4516758528</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M267.58,-48C289.9,-48 322.2,-48 347.21,-48\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"346.97,-51.5 356.97,-48 346.97,-44.5 346.97,-51.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"313.1\" y=\"-51.8\" font-family=\"Times,serif\" font-size=\"14.00\">&#45;1</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x10d3844f0>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "softmax.generate_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And use the `apply_fwd` method to calculate the softmax.  We can then verify that it is the same as our `staged_softmax` that we calculated earlier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "operations_softmax = softmax.apply_fwd()\n",
    "\n",
    "np.allclose(staged_softmax, operations_softmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, for the cool part.  We can use the same computational graph to do the backward pass!  Let's first visualize what the backward pass looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 7.0.6 (20230106.0513)\n",
       " -->\n",
       "<!-- Title: fwd_pass Pages: 1 -->\n",
       "<svg width=\"678pt\" height=\"139pt\"\n",
       " viewBox=\"0.00 0.00 678.09 139.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 135)\">\n",
       "<title>fwd_pass</title>\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-135 674.09,-135 674.09,4 -4,4\"/>\n",
       "<!-- 4516507376 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>4516507376</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"42.25\" cy=\"-95\" rx=\"42.49\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"42.25\" y=\"-91.3\" font-family=\"Times,serif\" font-size=\"14.00\">Multiply</text>\n",
       "</g>\n",
       "<!-- 4516759200 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>4516759200</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"527.09\" cy=\"-95\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"527.09\" y=\"-91.3\" font-family=\"Times,serif\" font-size=\"14.00\">Exp</text>\n",
       "</g>\n",
       "<!-- 4516507376&#45;&gt;4516759200 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>4516507376&#45;&gt;4516759200</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M82.03,-101.75C88.83,-102.69 95.85,-103.5 102.49,-104 270.72,-116.7 314.24,-120.99 482.09,-104 484.61,-103.74 487.2,-103.41 489.8,-103.02\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"490.36,-106.47 499.6,-101.29 489.15,-99.58 490.36,-106.47\"/>\n",
       "<text text-anchor=\"middle\" x=\"316.99\" y=\"-119.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(e^X)</text>\n",
       "</g>\n",
       "<!-- 4516758528 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>4516758528</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"236.49\" cy=\"-72\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"236.49\" y=\"-68.3\" font-family=\"Times,serif\" font-size=\"14.00\">Pow</text>\n",
       "</g>\n",
       "<!-- 4516507376&#45;&gt;4516758528 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>4516507376&#45;&gt;4516758528</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M83.55,-90.18C117.38,-86.13 165.33,-80.39 198.26,-76.45\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"198.65,-79.93 208.16,-75.27 197.82,-72.98 198.65,-79.93\"/>\n",
       "<text text-anchor=\"middle\" x=\"146.99\" y=\"-91.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(1 / sum(e^X))</text>\n",
       "</g>\n",
       "<!-- 4516758288 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>4516758288</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"643.09\" cy=\"-95\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"643.09\" y=\"-91.3\" font-family=\"Times,serif\" font-size=\"14.00\">X</text>\n",
       "</g>\n",
       "<!-- 4516759200&#45;&gt;4516758288 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>4516759200&#45;&gt;4516758288</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M554.53,-95C569.35,-95 588.08,-95 604.41,-95\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"604.28,-98.5 614.28,-95 604.28,-91.5 604.28,-98.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"585.09\" y=\"-98.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(X)</text>\n",
       "</g>\n",
       "<!-- 4516758192 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>4516758192</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"397.79\" cy=\"-72\" rx=\"27.1\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"397.79\" y=\"-68.3\" font-family=\"Times,serif\" font-size=\"14.00\">Sum</text>\n",
       "</g>\n",
       "<!-- 4516758528&#45;&gt;4516758192 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>4516758528&#45;&gt;4516758192</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M263.8,-72C289.74,-72 329.65,-72 358.99,-72\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"358.93,-75.5 368.93,-72 358.93,-68.5 358.93,-75.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"316.99\" y=\"-75.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(sum(e^X))</text>\n",
       "</g>\n",
       "<!-- 4516758240 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>4516758240</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"397.79\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"397.79\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">&#45;1</text>\n",
       "</g>\n",
       "<!-- 4516758528&#45;&gt;4516758240 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>4516758528&#45;&gt;4516758240</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M261,-64.03C287.82,-54.94 331.93,-39.99 362.57,-29.6\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"363.41,-33.01 371.76,-26.49 361.16,-26.38 363.41,-33.01\"/>\n",
       "</g>\n",
       "<!-- 4516758192&#45;&gt;4516759200 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>4516758192&#45;&gt;4516759200</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M424.24,-76.6C442.98,-79.98 468.72,-84.63 489.63,-88.41\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"488.72,-91.81 499.19,-90.14 489.97,-84.92 488.72,-91.81\"/>\n",
       "<text text-anchor=\"middle\" x=\"462.59\" y=\"-90.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(e^X)</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x10d346cb0>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "softmax.generate_graph(backward=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then run `apply_bwd` to run the backward pass.  We first call `zero_grad` to ensure that the initial gradients on all nodes are properly set to zero.\n",
    "\n",
    "Any parameter nodes with `needs_grad` set to `True` will now have a property called `grad`, where we can get the gradient.  We can get the gradient from `X`, and verify that it matches our staged gradient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "softmax.zero_grad()\n",
    "softmax.apply_bwd(loss_grad)\n",
    "\n",
    "\n",
    "operations_softmax_grad = X.grad\n",
    "np.allclose(staged_softmax_grad, operations_softmax_grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a bonus, we can also check the equations that we multiplied to calculate the partial derivative with respect to `X`.  You can see that we had to add together 2 different gradients, just like we did with the staged version.  The $\\partial$ symbol means partial derivative, and $\\frac{\\partial e^{X}}{\\partial X}$ means \"the partial derivative of $e^X$ with respect to $X$.  The way to interpret this partial derivative is \"as $X$ changes, how does $e^{X}$ change?\".  By multiplying the partial derivatives of each operation, we can get the larger partial derivative, which is $\\frac{\\partial L}{\\partial X}$ - how the loss changes as we change $X$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\frac{\\partial}{\\partial X} = \\frac{\\partial softmax(X)}{\\partial e^X}*\\frac{\\partial e^X}{\\partial X} + \\\\\\frac{\\partial softmax(X)}{\\partial 1 / sum(e^X)}*\\frac{\\partial 1 / sum(e^X)}{\\partial sum(e^X)}*\\frac{\\partial sum(e^X)}{\\partial e^X}*\\frac{\\partial e^X}{\\partial X}$"
      ],
      "text/plain": [
       "<IPython.core.display.Latex object>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "softmax.generate_derivative_chains()\n",
    "display_chain(X.display_partial_derivative())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We just built a computational graph of the softmax function, then ran a forward and backward pass.  This is exactly how deep learning frameworks like PyTorch and TensorFlow work.  They define common operators, like `torch.dot`, keep track of what operators you called, build a graph, and automatically run the backward pass based on the derivative of each operator.  You could say we built a miniature deep learning framework!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2-layer network\n",
    "\n",
    "Let's extend our framework to work across a 2-layer neural network.  We need to define 3 additional operations:\n",
    "\n",
    "- `MatMul` - to multiply two matrices.  The forward pass is `x@w`.\n",
    "- `Add` - add two values up.\n",
    "- `Relu` - a nonlinear activation function.  Anything below `0` will be set to `0`.\n",
    "\n",
    "Here are the operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "class MatMul(Node):\n",
    "    def forward(self, x, w):\n",
    "        return x @ w # multiply the two matrices.\n",
    "\n",
    "    def backward(self, grad):\n",
    "        x, w = self.cache\n",
    "        # return the input gradient times the weights as the gradient on x\n",
    "        # the input x values (from the forward pass) times the input gradient is the gradient on the weights\n",
    "        return grad @ w.T, x.T @ grad\n",
    "\n",
    "class Add(Node):\n",
    "    def forward(self, x, b):\n",
    "        return x + b\n",
    "\n",
    "    def backward(self, grad):\n",
    "        # Any change to x or b will scale the output the same amount\n",
    "        return grad, grad\n",
    "\n",
    "class Relu(Node):\n",
    "    def forward(self, x):\n",
    "        return np.maximum(x, 0)\n",
    "\n",
    "    def backward(self, grad):\n",
    "        x = self.cache[0]\n",
    "        new_grad = np.array(grad)\n",
    "        # The derivative of relu is 0 when the input\n",
    "        # in the forward pass was below 0\n",
    "        # 1 otherwise\n",
    "        new_grad[x < 0] = 0\n",
    "        return new_grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you train a neural network using a framework like PyTorch, you will usually be using the GPU.  GPUs enable us to parallelize operations and train neural networks much faster than we could with a CPU.\n",
    "\n",
    "One way to speed up GPU code is to fuse operators.  This means that we combine operators that run together into a single operator.  This combined operator can then be optimized into a kernel that runs better on the GPU.  We'll discuss this more in a future lesson.\n",
    "\n",
    "As an example of the process, below you can see the fusion of the `Softmax` operation.  The fused softmax used the softmax functions we defined before to do the forward and backward passes.\n",
    "\n",
    "We also fuse an entire neural network layer into the `Dense` operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "class Softmax(Node):\n",
    "    def forward(self, x):\n",
    "        return softmax_func(x)\n",
    "\n",
    "    def backward(self, grad):\n",
    "        x = self.cache[0]\n",
    "        softmax = self.forward(x)\n",
    "        return softmax_grad_func(softmax, grad)\n",
    "\n",
    "class Dense(Node):\n",
    "    def forward(self, x, w, b):\n",
    "        # Multiply by weight, add bias\n",
    "        return x @ w + b\n",
    "\n",
    "    def backward(self, grad):\n",
    "        x, w, b = self.cache\n",
    "        # Return 3 gradients for x, w, and b\n",
    "        return grad @ w.T, x.T @ grad, grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have the operators we need to define our neural network.  Let's first load in the data.  We'll use the same telescope data from an earlier lesson.  We have observations from a telescope, and we want to classify whether each observation is a star, galaxy, or quasar.\n",
    "\n",
    "We use a data wrapper that I wrote to load and split the data automatically.  We'll only use the training set in this lesson, but you can experiment with the other 2 sets as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "sys.path.append(os.path.abspath(\"../data\"))\n",
    "from csv_data import SkyServerDatasetWrapper\n",
    "\n",
    "# Load the data\n",
    "wrapper = SkyServerDatasetWrapper()\n",
    "[train_x, train_y], [valid_x, valid_y], [test_x, test_y] = wrapper.get_flat_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.16745842, -0.58492272,  1.03148637, -0.34855938, -0.83728027,\n",
       "        -0.94605772, -0.99534154, -0.83806089,  0.21085172, -0.21763043,\n",
       "        -0.36973112,  1.03148936,  1.30931064],\n",
       "       [ 0.16886159, -0.58311429,  0.05243046, -0.16653251, -0.15415531,\n",
       "        -0.08264457, -0.02604308, -0.83806089,  0.21085172, -0.21763043,\n",
       "        -0.36984929, -0.63621258, -0.87919741]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_x[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then need to one-hot encode our data so that we can use it for classification.  We create 3-element vectors where only one element is `1`, and the others are `0`.  The position of the `1` corresponds to the target we want to predict."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def encode(target, max_value=3):\n",
    "    # A matrix with 3 columns\n",
    "    encoded = np.zeros((target.shape[0], max_value))\n",
    "    # Setup the indices that we'll set to one\n",
    "    inds = (np.arange(0,target.shape[0]), target.reshape(-1))\n",
    "    # Set the target positions to 1\n",
    "    encoded[inds] = 1\n",
    "    return encoded\n",
    "\n",
    "train_y = encode(train_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1., 0., 0.],\n",
       "       [1., 0., 0.],\n",
       "       [0., 1., 0.],\n",
       "       [1., 0., 0.],\n",
       "       [1., 0., 0.]])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_y[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our training data, we can initialize our weights and biases.  We're taking in `13` features, so we'll setup our weights for the first layer accordingly.  We want to output `3` digits, one for each encoding position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Set a seed so we can reproduce results\n",
    "np.random.seed(0)\n",
    "w1 = np.random.rand(13, 10)\n",
    "b1 = np.random.rand(1, 10)\n",
    "w2 = np.random.rand(10, 3)\n",
    "b2 = np.random.rand(1, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now define our network in code!  This is very similar to networks we've built in the past.  We have our first layer, with a relu activation, then our second layer, then a softmax to get probabilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "X = Parameter(train_x, desc=\"X\", needs_grad=False)\n",
    "Y = Parameter(train_y, desc=\"y\", needs_grad=False)\n",
    "\n",
    "w1_param = Parameter(w1, desc=\"W1\")\n",
    "b1_param = Parameter(b1, desc=\"b1\")\n",
    "\n",
    "matmul1 = MatMul(X, w1_param, out=\"X @ W1\")\n",
    "add1 = Add(matmul1, b1_param, out=\"Z1\")\n",
    "\n",
    "layer1 = Relu(add1, out=\"A1\")\n",
    "\n",
    "w2_param = Parameter(w2, desc=\"W2\")\n",
    "b2_param = Parameter(b2, desc=\"b2\")\n",
    "matmul2 = MatMul(layer1, w2_param, out=\"Z1 @ W2\")\n",
    "add2 = Add(matmul2, b2_param, out=\"Z2\")\n",
    "\n",
    "softmax = Softmax(add2, out=\"softmax(Z2)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just like before, we can create a computation graph.  This one is more complex than before, but should still be readable!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 7.0.6 (20230106.0513)\n",
       " -->\n",
       "<!-- Title: fwd_pass Pages: 1 -->\n",
       "<svg width=\"862pt\" height=\"191pt\"\n",
       " viewBox=\"0.00 0.00 861.67 191.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 187)\">\n",
       "<title>fwd_pass</title>\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-187 857.67,-187 857.67,4 -4,4\"/>\n",
       "<!-- 4516503872 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>4516503872</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"812.73\" cy=\"-37\" rx=\"40.89\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"812.73\" y=\"-33.3\" font-family=\"Times,serif\" font-size=\"14.00\">Softmax</text>\n",
       "</g>\n",
       "<!-- 4516501952 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>4516501952</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"692.78\" cy=\"-37\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"692.78\" y=\"-33.3\" font-family=\"Times,serif\" font-size=\"14.00\">Add</text>\n",
       "</g>\n",
       "<!-- 4516501952&#45;&gt;4516503872 -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>4516501952&#45;&gt;4516503872</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M720.25,-37C732.03,-37 746.33,-37 760.08,-37\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"759.93,-40.5 769.93,-37 759.93,-33.5 759.93,-40.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"745.78\" y=\"-40.8\" font-family=\"Times,serif\" font-size=\"14.00\">Z2</text>\n",
       "</g>\n",
       "<!-- 4516511072 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>4516511072</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"534.48\" cy=\"-72\" rx=\"40.09\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"534.48\" y=\"-68.3\" font-family=\"Times,serif\" font-size=\"14.00\">MatMul</text>\n",
       "</g>\n",
       "<!-- 4516511072&#45;&gt;4516501952 -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>4516511072&#45;&gt;4516501952</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M570.95,-64.06C596.26,-58.39 630.18,-50.8 655.63,-45.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"656.2,-48.55 665.2,-42.95 654.67,-41.72 656.2,-48.55\"/>\n",
       "<text text-anchor=\"middle\" x=\"620.28\" y=\"-61.8\" font-family=\"Times,serif\" font-size=\"14.00\">Z1 @ W2</text>\n",
       "</g>\n",
       "<!-- 4516498688 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>4516498688</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"409.89\" cy=\"-103\" rx=\"27.1\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"409.89\" y=\"-99.3\" font-family=\"Times,serif\" font-size=\"14.00\">Relu</text>\n",
       "</g>\n",
       "<!-- 4516498688&#45;&gt;4516511072 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>4516498688&#45;&gt;4516511072</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M435.7,-96.73C450.69,-92.94 470.27,-87.99 487.96,-83.51\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"488.65,-86.95 497.49,-81.1 486.94,-80.16 488.65,-86.95\"/>\n",
       "<text text-anchor=\"middle\" x=\"465.69\" y=\"-94.8\" font-family=\"Times,serif\" font-size=\"14.00\">A1</text>\n",
       "</g>\n",
       "<!-- 4516498784 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>4516498784</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"303.59\" cy=\"-103\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"303.59\" y=\"-99.3\" font-family=\"Times,serif\" font-size=\"14.00\">Add</text>\n",
       "</g>\n",
       "<!-- 4516498784&#45;&gt;4516498688 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>4516498784&#45;&gt;4516498688</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M330.91,-103C343.14,-103 357.9,-103 371.28,-103\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"371.07,-106.5 381.07,-103 371.07,-99.5 371.07,-106.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"356.59\" y=\"-106.8\" font-family=\"Times,serif\" font-size=\"14.00\">Z1</text>\n",
       "</g>\n",
       "<!-- 4516497824 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>4516497824</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"151.3\" cy=\"-138\" rx=\"40.09\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"151.3\" y=\"-134.3\" font-family=\"Times,serif\" font-size=\"14.00\">MatMul</text>\n",
       "</g>\n",
       "<!-- 4516497824&#45;&gt;4516498784 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>4516497824&#45;&gt;4516498784</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M187.53,-129.79C211.39,-124.24 242.82,-116.92 266.79,-111.34\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"267.48,-114.77 276.43,-109.09 265.89,-107.95 267.48,-114.77\"/>\n",
       "<text text-anchor=\"middle\" x=\"234.09\" y=\"-126.8\" font-family=\"Times,serif\" font-size=\"14.00\">X @ W1</text>\n",
       "</g>\n",
       "<!-- 4516495904 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>4516495904</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-165\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-161.3\" font-family=\"Times,serif\" font-size=\"14.00\">X</text>\n",
       "</g>\n",
       "<!-- 4516495904&#45;&gt;4516497824 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>4516495904&#45;&gt;4516497824</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M53.04,-159.47C67.6,-156.26 86.42,-152.1 103.62,-148.31\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"104.13,-151.78 113.14,-146.2 102.62,-144.94 104.13,-151.78\"/>\n",
       "<text text-anchor=\"middle\" x=\"82.5\" y=\"-157.8\" font-family=\"Times,serif\" font-size=\"14.00\">X</text>\n",
       "</g>\n",
       "<!-- 4516509056 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>4516509056</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"27\" cy=\"-111\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"27\" y=\"-107.3\" font-family=\"Times,serif\" font-size=\"14.00\">W1</text>\n",
       "</g>\n",
       "<!-- 4516509056&#45;&gt;4516497824 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>4516509056&#45;&gt;4516497824</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M53.04,-116.53C67.6,-119.74 86.42,-123.9 103.62,-127.69\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"102.62,-131.06 113.14,-129.8 104.13,-124.22 102.62,-131.06\"/>\n",
       "<text text-anchor=\"middle\" x=\"82.5\" y=\"-128.8\" font-family=\"Times,serif\" font-size=\"14.00\">W1</text>\n",
       "</g>\n",
       "<!-- 4516497392 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>4516497392</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"151.3\" cy=\"-84\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"151.3\" y=\"-80.3\" font-family=\"Times,serif\" font-size=\"14.00\">b1</text>\n",
       "</g>\n",
       "<!-- 4516497392&#45;&gt;4516498784 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>4516497392&#45;&gt;4516498784</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M178.63,-85.21C200.21,-86.45 231.51,-88.8 258.59,-93 261.27,-93.42 264.03,-93.9 266.8,-94.43\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"266.05,-97.85 276.56,-96.48 267.49,-91 266.05,-97.85\"/>\n",
       "<text text-anchor=\"middle\" x=\"234.09\" y=\"-96.8\" font-family=\"Times,serif\" font-size=\"14.00\">b1</text>\n",
       "</g>\n",
       "<!-- 4516496672 -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>4516496672</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"409.89\" cy=\"-49\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"409.89\" y=\"-45.3\" font-family=\"Times,serif\" font-size=\"14.00\">W2</text>\n",
       "</g>\n",
       "<!-- 4516496672&#45;&gt;4516511072 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>4516496672&#45;&gt;4516511072</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M436.73,-51.94C448.77,-53.47 463.26,-55.54 476.19,-58 480.03,-58.73 483.99,-59.56 487.96,-60.43\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"487.06,-63.81 497.59,-62.64 488.63,-56.99 487.06,-63.81\"/>\n",
       "<text text-anchor=\"middle\" x=\"465.69\" y=\"-61.8\" font-family=\"Times,serif\" font-size=\"14.00\">W2</text>\n",
       "</g>\n",
       "<!-- 4516506896 -->\n",
       "<g id=\"node11\" class=\"node\">\n",
       "<title>4516506896</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"534.48\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"534.48\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">b2</text>\n",
       "</g>\n",
       "<!-- 4516506896&#45;&gt;4516501952 -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>4516506896&#45;&gt;4516501952</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M561.93,-19.05C584.76,-20.21 618.59,-22.54 647.78,-27 650.46,-27.41 653.22,-27.89 655.99,-28.42\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"655.25,-31.84 665.75,-30.46 656.68,-24.99 655.25,-31.84\"/>\n",
       "<text text-anchor=\"middle\" x=\"620.28\" y=\"-30.8\" font-family=\"Times,serif\" font-size=\"14.00\">b2</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x10d345e40>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "softmax.generate_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can create predictions using `apply_fwd`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.51408579, 0.29557176, 0.19034245],\n",
       "       [0.42508635, 0.26655849, 0.30835516]])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions = softmax.apply_fwd()\n",
    "\n",
    "predictions[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can run the backward pass of the network.  We'll first graph it out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 7.0.6 (20230106.0513)\n",
       " -->\n",
       "<!-- Title: fwd_pass Pages: 1 -->\n",
       "<svg width=\"955pt\" height=\"191pt\"\n",
       " viewBox=\"0.00 0.00 954.67 191.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 187)\">\n",
       "<title>fwd_pass</title>\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-187 950.67,-187 950.67,4 -4,4\"/>\n",
       "<!-- 4516503872 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>4516503872</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"40.95\" cy=\"-37\" rx=\"40.89\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"40.95\" y=\"-33.3\" font-family=\"Times,serif\" font-size=\"14.00\">Softmax</text>\n",
       "</g>\n",
       "<!-- 4516501952 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>4516501952</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"175.89\" cy=\"-37\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"175.89\" y=\"-33.3\" font-family=\"Times,serif\" font-size=\"14.00\">Add</text>\n",
       "</g>\n",
       "<!-- 4516503872&#45;&gt;4516501952 -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>4516503872&#45;&gt;4516501952</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M82.2,-37C99.71,-37 120.11,-37 137.31,-37\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"137.11,-40.5 147.11,-37 137.11,-33.5 137.11,-40.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"115.39\" y=\"-40.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(Z2)</text>\n",
       "</g>\n",
       "<!-- 4516511072 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>4516511072</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"349.19\" cy=\"-72\" rx=\"40.09\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"349.19\" y=\"-68.3\" font-family=\"Times,serif\" font-size=\"14.00\">MatMul</text>\n",
       "</g>\n",
       "<!-- 4516501952&#45;&gt;4516511072 -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>4516501952&#45;&gt;4516511072</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M202.17,-42.16C228.14,-47.47 269.19,-55.86 301.16,-62.39\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"300.06,-65.74 310.56,-64.31 301.46,-58.88 300.06,-65.74\"/>\n",
       "<text text-anchor=\"middle\" x=\"255.89\" y=\"-62.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(Z1 @ W2)</text>\n",
       "</g>\n",
       "<!-- 4516506896 -->\n",
       "<g id=\"node11\" class=\"node\">\n",
       "<title>4516506896</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"349.19\" cy=\"-18\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"349.19\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">b2</text>\n",
       "</g>\n",
       "<!-- 4516501952&#45;&gt;4516506896 -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>4516501952&#45;&gt;4516506896</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M201.43,-30.72C207.74,-29.31 214.54,-27.94 220.89,-27 250.79,-22.57 284.96,-20.31 310.37,-19.16\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"310.37,-22.66 320.22,-18.75 310.08,-15.67 310.37,-22.66\"/>\n",
       "<text text-anchor=\"middle\" x=\"255.89\" y=\"-30.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(b2)</text>\n",
       "</g>\n",
       "<!-- 4516498688 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>4516498688</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"489.78\" cy=\"-103\" rx=\"27.1\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"489.78\" y=\"-99.3\" font-family=\"Times,serif\" font-size=\"14.00\">Relu</text>\n",
       "</g>\n",
       "<!-- 4516511072&#45;&gt;4516498688 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>4516511072&#45;&gt;4516498688</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M385.52,-79.9C406,-84.48 431.73,-90.24 452.38,-94.86\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"451.57,-98.26 462.1,-97.03 453.1,-91.43 451.57,-98.26\"/>\n",
       "<text text-anchor=\"middle\" x=\"425.99\" y=\"-95.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(A1)</text>\n",
       "</g>\n",
       "<!-- 4516496672 -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>4516496672</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"489.78\" cy=\"-49\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"489.78\" y=\"-45.3\" font-family=\"Times,serif\" font-size=\"14.00\">W2</text>\n",
       "</g>\n",
       "<!-- 4516511072&#45;&gt;4516496672 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>4516511072&#45;&gt;4516496672</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M384.27,-62.91C391.87,-61.09 399.92,-59.34 407.49,-58 421.8,-55.47 437.66,-53.52 451.54,-52.1\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"451.54,-55.62 461.16,-51.17 450.87,-48.65 451.54,-55.62\"/>\n",
       "<text text-anchor=\"middle\" x=\"425.99\" y=\"-61.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(W2)</text>\n",
       "</g>\n",
       "<!-- 4516498784 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>4516498784</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"611.08\" cy=\"-103\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"611.08\" y=\"-99.3\" font-family=\"Times,serif\" font-size=\"14.00\">Add</text>\n",
       "</g>\n",
       "<!-- 4516498688&#45;&gt;4516498784 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>4516498688&#45;&gt;4516498784</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M517.26,-103C533.4,-103 554.29,-103 572.17,-103\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"572.17,-106.5 582.17,-103 572.17,-99.5 572.17,-106.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"550.58\" y=\"-106.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(Z1)</text>\n",
       "</g>\n",
       "<!-- 4516497824 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>4516497824</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"779.38\" cy=\"-138\" rx=\"40.09\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"779.38\" y=\"-134.3\" font-family=\"Times,serif\" font-size=\"14.00\">MatMul</text>\n",
       "</g>\n",
       "<!-- 4516498784&#45;&gt;4516497824 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>4516498784&#45;&gt;4516497824</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M637.34,-108.32C662.15,-113.54 700.62,-121.64 731.12,-128.05\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"730.39,-131.48 740.9,-130.11 731.83,-124.63 730.39,-131.48\"/>\n",
       "<text text-anchor=\"middle\" x=\"688.58\" y=\"-128.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(X @ W1)</text>\n",
       "</g>\n",
       "<!-- 4516497392 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>4516497392</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"779.38\" cy=\"-84\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"779.38\" y=\"-80.3\" font-family=\"Times,serif\" font-size=\"14.00\">b1</text>\n",
       "</g>\n",
       "<!-- 4516498784&#45;&gt;4516497392 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>4516498784&#45;&gt;4516497392</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M636.62,-96.74C642.93,-95.32 649.73,-93.95 656.08,-93 684.31,-88.78 716.51,-86.51 740.78,-85.31\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"740.67,-88.82 750.5,-84.87 740.35,-81.82 740.67,-88.82\"/>\n",
       "<text text-anchor=\"middle\" x=\"688.58\" y=\"-96.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(b1)</text>\n",
       "</g>\n",
       "<!-- 4516495904 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>4516495904</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"919.67\" cy=\"-165\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"919.67\" y=\"-161.3\" font-family=\"Times,serif\" font-size=\"14.00\">X</text>\n",
       "</g>\n",
       "<!-- 4516497824&#45;&gt;4516495904 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>4516497824&#45;&gt;4516495904</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M816.71,-145.09C836.92,-149.04 861.99,-153.93 882.21,-157.88\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"881.51,-161.31 892,-159.79 882.85,-154.44 881.51,-161.31\"/>\n",
       "</g>\n",
       "<!-- 4516509056 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>4516509056</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"919.67\" cy=\"-111\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"919.67\" y=\"-107.3\" font-family=\"Times,serif\" font-size=\"14.00\">W1</text>\n",
       "</g>\n",
       "<!-- 4516497824&#45;&gt;4516509056 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>4516497824&#45;&gt;4516509056</title>\n",
       "<path fill=\"none\" stroke=\"red\" d=\"M816.71,-130.91C836.92,-126.96 861.99,-122.07 882.21,-118.12\"/>\n",
       "<polygon fill=\"red\" stroke=\"red\" points=\"882.85,-121.56 892,-116.21 881.51,-114.69 882.85,-121.56\"/>\n",
       "<text text-anchor=\"middle\" x=\"856.17\" y=\"-129.8\" font-family=\"Times,serif\" font-size=\"14.00\">d(W1)</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x152e57dc0>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "softmax.generate_graph(backward=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we'll run our backward pass.  We first calculate our loss gradient, then pass it into our backward pass as the incoming gradient:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "softmax.zero_grad() # zero out the gradients\n",
    "loss_grad = nll_grad(train_y, predictions) # compute the loss gradient\n",
    "softmax.apply_bwd(loss_grad) # pass the loss gradient into the softmax, and run the backward pass\n",
    "softmax.generate_derivative_chains() # generate our derivative equations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each of our parameters should now have a gradient.  We'll only show `w2`, but you can look at the other parameters as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1078.86623385,  1452.92238635, -2531.7886202 ],\n",
       "       [  685.34169939,  1438.06439063, -2123.40609002],\n",
       "       [   11.39607931,  3488.94754329, -3500.3436226 ],\n",
       "       [ 1990.1147882 ,   471.22083925, -2461.33562744],\n",
       "       [ -538.48614737,  1989.95610171, -1451.46995434]])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w2_param.grad[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also look at the partial derivative chains to see how we calculate the derivative with respect to a specific parameter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\frac{\\partial}{\\partial W2} = \\frac{\\partial softmax(Z2)}{\\partial Z2}*\\frac{\\partial Z2}{\\partial Z1 @ W2}*\\frac{\\partial Z1 @ W2}{\\partial W2}$"
      ],
      "text/plain": [
       "<IPython.core.display.Latex object>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "display_chain(w2_param.display_partial_derivative())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\frac{\\partial}{\\partial W1} = \\frac{\\partial softmax(Z2)}{\\partial Z2}*\\frac{\\partial Z2}{\\partial Z1 @ W2}*\\frac{\\partial Z1 @ W2}{\\partial A1}*\\frac{\\partial A1}{\\partial Z1}*\\frac{\\partial Z1}{\\partial X @ W1}*\\frac{\\partial X @ W1}{\\partial W1}$"
      ],
      "text/plain": [
       "<IPython.core.display.Latex object>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "display_chain(w1_param.display_partial_derivative())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We just built a computational graph, and used it to do the full forward and backward pass for a neural network!  If you want, you can extend this to update the parameters and train the network.  You would just need to set a learning rate, then subtract the gradient from each parameter.  You would have to set a batch size, and iterate through the data as well.\n",
    "\n",
    "This has hopefully given you a good look at how backpropagation, works, and how we compute the partial derivatives of each operation, then multiply them out.\n",
    "\n",
    "Let's do a quick verification to make sure that we did everything correctly.  We can implement the network forward and backward pass like we did in an earlier lesson:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Forward pass\n",
    "l1 = train_x @ w1 + b1\n",
    "l1_activated = np.maximum(l1, 0)\n",
    "l2 = l1_activated @ w2 + b2\n",
    "probs = softmax_func(l2)\n",
    "\n",
    "# Loss\n",
    "loss_grad = nll_grad(train_y, probs)\n",
    "\n",
    "# L2 gradients\n",
    "sm_grad = softmax_grad_func(probs, loss_grad)\n",
    "l2_w_grad = l1_activated.T @ sm_grad\n",
    "l2_b_grad = sm_grad.sum(axis=0)\n",
    "\n",
    "# L1 gradients\n",
    "l1_grad = sm_grad @ w2.T\n",
    "l1_grad[l1 < 0] = 0\n",
    "l1_w_grad = train_x.T @ l1_grad\n",
    "l1_b_grad = l1_grad.sum(axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can verify that our computational graph matches the manual results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(l1_w_grad, w1_param.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wrap-up\n",
    "\n",
    "We did a lot in this lesson!  We learned how to break apart a derivative into steps, then compute each step separately.  Then, we constructed a computational graph and ran the forward and backward passes.\n",
    "\n",
    "I recommend doing some experimentation with the graph, and making sure you really understand how everything is working.  In the next lesson, we'll use PyTorch to automatically construct the graph for us."
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}