Run backpropagation in the right order
Every Value below already knows its local gradients: calling v._backward() passes v.grad on to the values v was made from (in v._prev). What is missing is backward(out), which runs them all.
Set out.grad to 1 (out changes by 1 per unit of itself). Then call _backward() on every value in the graph, starting at out and working back to the leaves, so a value has received all of its gradient before it passes it on.
A topological order does it: visit a value’s inputs before the value itself (a depth-first search over _prev, remembering what you’ve seen), then go through that list backwards. 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():
self.grad += out.grad
other.grad += out.grad
out._backward = _backward
return out
def __mul__(self, other):
out = Value(self.data * other.data, (self, other))
def _backward():
self.grad += other.data * out.grad
other.grad += self.data * out.grad
out._backward = _backward
return out
def tanh(self):
t = math.tanh(self.data)
out = Value(t, (self,))
def _backward():
self.grad += (1 - t * t) * out.grad
out._backward = _backward
return out
def backward(out):
# Set out.grad to 1, then call _backward() on every value,
# starting from out and working back towards the leaves.
pass
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