Build a Value Object in Python That Remembers Where It Came From

The idea

Write z = 5 * 4 in Python and you get 20. That is all you get. The number 20 does not know it came from a multiplication, and it does not know that 5 and 4 were involved.

Automatic differentiation needs that missing information. To compute derivatives you have to walk back through every step that produced the result. So instead of a bare number, we store each number in a small object that also keeps two extra things: which values it was made from, and which operation made it.

The code

This is the whole thing. Read it once, then we will take it apart line by line.

1class Value:
2    def __init__(self, data, _children=(), _op=''):
3        self.data = data
4        self._prev = set(_children)
5        self._op = _op
6
7    def __add__(self, other):
8        other = other if isinstance(other, Value) else Value(other)
9        return Value(self.data + other.data, (self, other), '+')
10
11    def __mul__(self, other):
12        other = other if isinstance(other, Value) else Value(other)
13        return Value(self.data * other.data, (self, other), '*')

__add__ and __mul__ are Python's hooks for the + and * operators. Defining them means you can write a + b and Python will call a.__add__(b) for you. That is what “operator overloading” means here.

_prev, _op, self and other

data is the plain number, e.g. 2.

_prev is a Python set of Value objects: the previous values this one was computed from. We build it with set(_children) because a set stores each input only once — if you write x + x, the same object appears twice in the tuple but only once in the set. For an input you create yourself,_children defaults to the empty tuple (), so _prev is an empty set.

_op is a short string naming the operation that produced this value: '+', '*', or '' (empty) if nothing produced it.

Inside a method, self is the object on the left of the operator and other is the one on the right. In a + b, self is a and other is b.

The line other = other if isinstance(other, Value) else Value(other) just lets you write a + 3: a raw number gets wrapped in a Value first.

Finally, (self, other) is the tuple we hand to the new Value as its _children. It says: these two values are what produced this result.

Step by step

Run these six lines:

1a = Value(2)
2b = Value(3)
3c = Value(4)
4
5sum = a + b     # calls a.__add__(b)
6z = sum * c     # calls sum.__mul__(c)

a, b and c are inputs: empty _prev, empty _op.

For sum = a + b, self is a and other is b. A brand new Value is returned with data 2 + 3, children (a, b) and op '+'.

For z = sum * c, self is sum and other is c. Note that a and b are not listed on z directly — they are reachable through sum.

Here is exactly what the objects hold afterwards:

1sum.data   # 5
2sum._prev  # {a, b}
3sum._op    # '+'
4
5z.data     # 20
6z._prev    # {sum, c}
7z._op      # '*'

The graph

Those _prev links form a small graph. Arrows point from inputs to the result they produced:

a ──┐
    ├─ + ──> sum ──┐
b ──┘              ├─ * ──> z
c ─────────────────┘

Each edge is one _prev reference, and each operation node is one _op string. Starting at z you can follow the links backwards and reach every value that contributed to it.

The graph never loops back on itself, because every operation returns a brand new object rather than modifying an existing one — a value can never end up as one of its own inputs.

No gradients are calculated yet. At this stage we only create Value objects and record the computation graph. That recorded graph is what makes backpropagation possible later, once each node also carries a gradient.

Try it

The same graph, in the browser. Change a, b or c and click a node to see where its value came from. No Python runs here.

What we built

• A Value stores a single number in data.
• + and * return new Value objects instead of plain numbers.
• _prev and _op record how each result was made.