Give each operation its local gradient
Below is the core of an autograd engine like micrograd. A Value remembers the values it was made from, and out._backward must pass out.grad back to them, each scaled by its local gradient.
Fill in the three _backward bodies: + (both inputs get out.grad), × (each input gets the other input’s data × out.grad) and tanh (its input gets (1 − t²) × out.grad). Always add to grad with +=: a value used twice collects a gradient from both uses.
grads_of(expr, vals, grad) (gradsOf in JavaScript and TypeScript) is written for you: it builds an expression from vals, runs backward, and copies each leaf’s gradient into grad. For neuron it is tanh(x1·w1 + x2·w2 + b); for reuse it is a·b + a·a (a is used three times); for chain it is tanh applied to out × w ten times over, starting from x.
import math
class Value:
def __init__(self, data, children=()):
self.data = data
self.grad = 0.0
self._prev = children
self._backward = lambda: None
def __add__(self, other):
out = Value(self.data + other.data, (self, other))
def _backward():
pass # TODO
out._backward = _backward
return out
def __mul__(self, other):
out = Value(self.data * other.data, (self, other))
def _backward():
pass # TODO
out._backward = _backward
return out
def tanh(self):
t = math.tanh(self.data)
out = Value(t, (self,))
def _backward():
pass # TODO
out._backward = _backward
return out
def backward(out):
order = []
seen = set()
def visit(v):
if v not in seen:
seen.add(v)
for child in v._prev:
visit(child)
order.append(v) # after everything it was made from
visit(out)
out.grad = 1.0
for v in reversed(order):
v._backward()
def build(expr, vals):
leaves = [Value(v) for v in vals]
if expr == 'neuron':
x1, w1, x2, w2, b = leaves
out = (x1 * w1 + x2 * w2 + b).tanh()
elif expr == 'reuse':
a, b = leaves
out = a * b + a * a
else:
x, w = leaves
out = x
for _ in range(10):
out = (out * w).tanh()
return leaves, out
def grads_of(expr, vals, grad):
leaves, out = build(expr, vals)
backward(out)
for i in range(len(leaves)):
grad[i] = round(leaves[i].grad, 9) # same digits in every language
return grad