A Scalar Autograd Engine

The problem

In PyTorch you compute a loss and call loss.backward(), and every parameter gets its gradient. Behind that one call there are three steps:

1. during the forward pass, remember every operation and its inputs (the computational graph);
2. sort the operations so that each one comes after all the operations that use its result;
3. walk them in that order, applying each operation's local rule from the chain rule.

Here you build that engine for single numbers, in the style of Karpathy's micrograd. Real frameworks work the same way on whole arrays.

The Value class

A Value wraps one number and keeps three extra things:

• data: the number itself;
• grad: the gradient of the final result with respect to this value, starting at 0.0;
• _prev: the Values it was computed from, and _backward: a function that sends this value's gradient to them.

a, b = Value(2.0), Value(-3.0)
c = a * b + a   # c.data: -4.0
c.backward()
a.grad, b.grad  # (-2.0, 2.0)

The gradients are right: $c = ab + a$, so $\partial c / \partial a = b + 1 = -2$ and $\partial c / \partial b = a = 2$.

Step 1: remember the graph

Every operation returns a new Value whose _prev holds its inputs. In the example, a * b creates a value d with _prev = (a, b) and data = -6, then d + a creates c with _prev = (d, a) and data = -4:

a ──┐
    ├─ (*) ─ d ──┐
b ──┘            ├─ (+) ─ c
a ───────────────┘

The new value also gets a _backward function with the local rule of its operation. When it runs, it reads the value's own .grad (the upstream gradient) and adds each input's share to that input's .grad:

Step 2: sort the graph

c._backward must run before d._backward: d only knows its gradient once c has sent it. The rule is: every value comes after all the values that use it. A depth-first search builds the opposite order, every value after its inputs, and reversing it gives the order we want:

1. to visit a value: if it was already visited, stop; otherwise mark it visited, visit each of its _prev, then append the value to a list;
2. start the visit from the final value, c;
3. reverse the list.

Example. Visiting c visits d, which visits a (appended) and b (appended), then d is appended; back in c, a was already visited; finally c is appended. The list is [a, b, d, c], reversed [c, d, b, a]: c comes before d, and d before a and b.

Step 3: walk the graph

Set c.grad = 1 (the derivative of anything with respect to itself is 1), then call _backward on every value in the sorted order.

Example. In the order [c, d, b, a]:

1. c is d + a: it adds 1 to d.grad and 1 to a.grad. Now d.grad = 1, a.grad = 1.
2. d is a * b: it adds b.data * d.grad = -3 to a.grad and a.data * d.grad = 2 to b.grad. Now a.grad = 1 - 3 = -2, b.grad = 2.
3. b and a were not computed from anything: their _backward does nothing.

Final result: a.grad = -2.0, b.grad = 2.0. Because a is used twice, its two contributions, $1$ and $-3$, were added.

What to implement

The four methods marked TODO:

• __add__ and __mul__: compute the result, create the new Value with its inputs in _prev, and give it a _backward with the rule of the table. The other operand can be a plain number such as 2: wrap it in a Value first.
• tanh(): the same pattern, with math.tanh and the rule of the table.
• backward(): sort the graph (step 2), set this value's .grad to 1, and call every _backward in order (step 3).

__radd__, __rmul__, __neg__ and __sub__ are given: they are built from your __add__ and __mul__, which is how frameworks get many operations out of a few backward rules.

Requirements

• Add, never assign: a * a must give a.grad == 2 * a.data. A value used twice receives two contributions, and both must count.
• Each value's _backward runs once, even if several paths reach it: that is what the visited set is for.
• Every .grad starts at 0. You don't need to support calling backward() twice.

Hint

Define _backward as a function inside __mul__, after creating out. It can read self, other and out, and it reads out.grad when it runs (during step 3), not when it is defined. Then attach it with out._backward = _backward.