Scalar autograd: PyTorch, from scratch, then a fused sigmoid

We will start with one graph:

\[ m=wx,\qquad a=m+b,\qquad e=a-y,\qquad L=e^2 \]

with \(w=2\), \(x=3\), \(b=1\), and \(y=10\).

First we let PyTorch differentiate it. Then we build the smallest useful version of the same idea ourselves. The point is not to replace PyTorch—it is to see what .backward() does. A final optional example then compares one fused sigmoid operation with the same sigmoid expanded into atomic operations.

Complete scalar computation graph with every forward value and loss gradient

Forward computes left → right. Backward begins at L.grad = 1 and travels right → left. The teal number in each box is that node’s final ∂L/∂node. On a phone, scroll sideways.

1 · The whole example in PyTorch

Each number below is a scalar tensor. We set requires_grad=True because we want PyTorch to calculate its loss derivative. In ordinary training, the target y would normally be fixed; here we track it only so the output matches our complete paper calculation.

PyTorch keeps .grad automatically for the four leaf tensors. retain_grad() asks it to keep gradients for the intermediate values too.

import torch

torch.set_default_dtype(torch.float64)

w = torch.tensor(2.0, requires_grad=True)
x = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
y = torch.tensor(10.0, requires_grad=True)

# Forward pass
m = w * x
a = m + b
e = a - y
L = e ** 2

for node in (m, a, e, L):
    node.retain_grad()

print("Forward: m =", m.item(), ", a =", a.item(),
      ", e =", e.item(), ", L =", L.item())
Forward: m = 6.0 , a = 7.0 , e = -3.0 , L = 9.0

Now the important line:

L.backward()

print("w.grad =", w.grad)
print("x.grad =", x.grad)
print("b.grad =", b.grad)
print("y.grad =", y.grad)
w.grad = tensor(-18.)
x.grad = tensor(-12.)
b.grad = tensor(-6.)
y.grad = tensor(6.)

That is PyTorch autograd. The forward pass created a graph; .backward() sent a seed gradient of \(1\) from \(L\) through that graph in reverse.

For comparison with the paper calculation, here is every stored value:

torch_nodes = {"w": w, "x": x, "m": m, "b": b,
               "a": a, "y": y, "e": e, "L": L}
torch_reference = {
    name: (node.item(), node.grad.item())
    for name, node in torch_nodes.items()
}

print(f"{'node':<5} {'value':>8} {'grad':>8}")
for name, (value, grad) in torch_reference.items():
    print(f"{name:<5} {value:8.1f} {grad:8.1f}")

expected = {
    "w": (2, -18), "x": (3, -12), "m": (6, -6), "b": (1, -6),
    "a": (7, -6), "y": (10, 6), "e": (-3, -6), "L": (9, 1),
}
assert torch_reference == expected
node     value     grad
w          2.0    -18.0
x          3.0    -12.0
m          6.0     -6.0
b          1.0     -6.0
a          7.0     -6.0
y         10.0      6.0
e         -3.0     -6.0
L          9.0      1.0

2 · A tiny autograd engine from scratch

Suppose one operation \(f\) takes a value \(u\) and produces \(v=f(u)\). In this notebook’s convention, the forward construction makes \(u\) a direct parent (or operand) of \(v\), while \(v\) is the output (or child) created from \(u\). These words describe one operation in the computation graph—not a whole neural-network layer. Libraries sometimes choose different names; here, follow v.parents, whose links point back to the direct operands.

Parent and child in a computation graph and the storage used for one backward update

v.parents owns the ParentLink, and that link points back to u. The engine does not need u to keep a list of its children. On a phone, scroll sideways.

Concrete link from our graph: in m = w × x, choose v = m and the parent u = w. Backward reads upstream m.grad = −6, reads local ∂m/∂w = x = 3 from that link, computes the temporary contribution −6 × 3 = −18, and adds it to w.grad. The other parent link uses u = x, local derivative w = 2, and contributes −12 to x.grad.

For the autograd calculation, a Value needs only:

  • its number in data,
  • its accumulated loss-gradient buffer in grad, and
  • ordered links to the operands that directly produced it.

Our teaching class also stores label and op so diagrams can say “m” and “×”. They are display metadata: changing those strings does not change the forward number or any gradient.

Each ParentLink has exactly two fields: .value points to the parent operand, and .local_grad holds the evaluated local derivative along that edge. For example, \(m=wx\) remembers \((w,\partial m/\partial w=x)\) and \((x,\partial m/\partial x=w)\).

The gradient names are always relative to the operation currently running:

  • upstream: \(g_v=\partial L/\partial v\), already accumulated in v.grad;
  • local: \(\partial v/\partial u\), stored in the link from output \(v\) to parent \(u\);
  • edge contribution to the parent (the “downstream contribution” in our color legend): \(\Delta g_u=g_v(\partial v/\partial u)\), computed during backward and added to u.grad.

The edge contribution is temporary: only the accumulated result in u.grad remains. w is a leaf because it has no dependencies. L is the forward output or sink; backward(L) treats it as the starting node—the root of the reverse traversal.

Before writing backward: decide when a node is ready

A node must not send its gradient to its parents until all gradient contributions arriving at that node have been added to its .grad buffer. We therefore need a dependency-safe processing order. Start at L, follow its saved ParentLinks toward the inputs, and append each node only after all its parents have been appended. Reversing the resulting list gives the safe order for backward.

Dependency-safe order for forward values, its reverse for backward, and an example of a wrong order

The formal name for any ordering that puts every dependency before the value that uses it is a topological order. We first understand the readiness rule; the name is secondary. On a phone, scroll sideways.

First apply the append-after-parents rule to one small part of the graph:

visit(m):
  visit(w) → w has no parents → append w
  visit(x) → x has no parents → append x
  both parents are ready       → append m

Starting from L applies that same rule recursively to the whole graph. With our stored parent order, the traversal produces exactly

dependency-first:  w, x, m, b, a, y, e, L
process back:  L, e, y, a, b, m, x, w

Other dependency-safe lists are possible—for example, two independent leaves can swap places. seen matters when a value feeds several later operations: following links from L may reach that same object more than once, but it must be appended and processed only once. seen deduplicates Value objects in the node schedule—not edges. If an operation is w * w, it still stores two ParentLinks, and backward still processes both contributions. For \(r=w^2\) at \(w=2\), w appears once in the node schedule, but the two links each contribute \(2\); w.grad += contribution therefore gives \(4\).

class ParentLink:
    def __init__(self, value, local_grad):
        self.value = value
        self.local_grad = float(local_grad)

class Value:
    def __init__(self, data, label="", parents=(), op=""):
        self.data = float(data)
        self.grad = 0.0
        self.label = label
        self.parents = tuple(parents)
        self.op = op

def multiply(u, v, label):
    return Value(u.data * v.data, label,
                 parents=(ParentLink(u, v.data),
                          ParentLink(v, u.data)), op="×")

def add(u, v, label):
    return Value(u.data + v.data, label,
                 parents=(ParentLink(u, 1.0),
                          ParentLink(v, 1.0)), op="+")

def subtract(u, v, label):
    return Value(u.data - v.data, label,
                 parents=(ParentLink(u, 1.0),
                          ParentLink(v, -1.0)), op="−")

def square(u, label):
    return Value(u.data ** 2, label,
                 parents=(ParentLink(u, 2 * u.data),), op="²")

def dependency_safe_order(root):
    """Return reachable Values with every parent before its output."""
    safe_order = []
    seen = set()

    def append_after_parents(node):
        # A shared Value may be reachable from the loss by several paths.
        # Visit and append that object only once.
        if id(node) in seen:
            return
        seen.add(id(node))

        # Follow output -> ParentLink -> operand, starting from the loss.
        for link in node.parents:
            append_after_parents(link.value)

        # Only now are all direct parents earlier in safe_order.
        safe_order.append(node)

    append_after_parents(root)
    return safe_order

def backward(root):
    safe_order = dependency_safe_order(root)

    # 1. Clear old accumulated gradients, then seed the loss.
    for node in safe_order:
        node.grad = 0.0
    root.grad = 1.0

    # 2. Reverse the safe order. A node's full upstream gradient is
    # ready before that node sends contributions to its parents.
    steps = []
    for output in reversed(safe_order):
        # One saved parent link gives one chain-rule update.
        for link in output.parents:
            parent = link.value
            upstream = output.grad
            local = link.local_grad
            contribution = upstream * local
            before = parent.grad
            parent.grad += contribution

            # Keep a teaching trace; autograd only needs the update above.
            steps.append({
                "output": output.label,
                "upstream": upstream,
                "parent": parent.label,
                "local": local,
                "downstream": contribution,
                "before": before,
                "after": parent.grad,
            })
    return steps

The visible code separates finding a safe order from doing the calculus:

  1. dependency_safe_order(root) starts at L. append_after_parents follows each stored link to a direct operand and calls itself there first. Only after those calls return does it append the current node. seen makes a shared object a no-op on its second visit.
  2. backward(root) clears every reachable .grad, then seeds L.grad = 1 because \(\partial L/\partial L=1\).
  3. reversed(safe_order) processes L, e, y, a, b, m, x, w. At every saved link from output \(v\) to parent \(u\), it reads the now-complete upstream gradient from v.grad, multiplies by the saved local derivative, and accumulates the result in u.grad.

Leaves such as w still appear in the processing list. They simply have no parent links, so there is nothing further to update when their turn arrives.

quantity where it lives
upstream \(g_v\) already accumulated in v.grad
local \(\partial v/\partial u\) saved in link.local_grad during the forward pass
edge contribution to parent \(\Delta g_u\) temporary variable contribution for this one edge
accumulated \(g_u\) updated in parent.grad

The autograd graph does not store a separate downstream gradient forever. It computes one edge contribution, adds it to the parent’s buffer, and that buffer later becomes the upstream gradient for the parent. Our returned steps list is only a teaching log: it copies each contribution and the before/after values so we can display them.

Three deliberate boundaries keep this engine small:

  • it assumes an acyclic computation graph (a DAG);
  • it seeds a scalar loss with 1; vector outputs would need an explicit upstream seed;
  • it clears reachable .grad buffers at the start of every call. PyTorch normally accumulates gradients across .backward() calls until you clear them.

Build the same forward graph, one readable line per operation. The interactive trace immediately below is generated from these actual Value objects and their stored ParentLinks—not from a separate hand-written event list. On a phone, scroll the graph sideways:

sw = Value(2.0, label="w")
sx = Value(3.0, label="x")
sb = Value(1.0, label="b")
sy = Value(10.0, label="y")

sm = multiply(sw, sx, "m")
sa = add(sm, sb, "a")
se = subtract(sa, sy, "e")
sL = square(se, "L")

print("What m=wx stored during forward:")
for link in sm.parents:
    print(f"  parent {link.value.label}: local ∂m/∂{link.value.label} = {link.local_grad:g}")

safe_order = dependency_safe_order(sL)
print("\nDependency-first order:", " → ".join(node.label for node in safe_order))
print("Backward will process:   ", " → ".join(node.label for node in reversed(safe_order)))

display(draw_graph(sL, show_grad=False))

topology_events = trace_dependency_safe_order(sL)
show_topological_sort_animation(sL, topology_events)
What m=wx stored during forward:
  parent w: local ∂m/∂w = 3
  parent x: local ∂m/∂x = 2

Dependency-first order: w → x → m → b → a → y → e → L
Backward will process:    L → e → y → a → b → m → x → w
value_w w value 2 grad — op_m × value_w->op_m ∂m/∂w=3 value_x x value 3 grad — value_x->op_m ∂m/∂x=2 value_m m value 6 grad — op_a + value_m->op_a ∂a/∂m=1 op_m->value_m value_b b value 1 grad — value_b->op_a ∂a/∂b=1 value_a a value 7 grad — op_e − value_a->op_e ∂e/∂a=1 op_a->value_a value_y y value 10 grad — value_y->op_e ∂e/∂y=-1 value_e e value -3 grad — op_L ² value_e->op_L ∂L/∂e=-6 op_e->value_e value_L L value 9 grad — op_L->value_L
Dependency-safe ordering, step by step
Uses the live graph and ParentLinks built above.
Open full screen

If the embedded view does not load, open the interactive in a new tab.

Dependency-safe order: w → x → m → b → a → y → e → L
Backward schedule: L → e → y → a → b → m → x → w

Now run backward once, then inspect the result at three levels:

  1. the edge-by-edge trace shows every chain-rule multiplication and accumulation;
  2. the compact graph shows the whole computation without overcrowding it;
  3. the complete state cards expose every field on every Value, including every saved parent link.

Only steps = backward(sL) performs differentiation. The two show_... helpers and draw_graph are teaching displays; removing them would not change any gradient.

steps = backward(sL)
show_backward_steps(steps)

display(draw_graph(sL, show_grad=True))
show_complete_state(sL)
Seed: L.grad = ∂L/∂L = 1
1. L → e
gL = 1  ×  ∂L/∂e = -6  =  Δge = -6
e.grad: 0 → -6
2. e → a
ge = -6  ×  ∂e/∂a = 1  =  Δga = -6
a.grad: 0 → -6
3. e → y
ge = -6  ×  ∂e/∂y = -1  =  Δgy = 6
y.grad: 0 → 6
4. a → m
ga = -6  ×  ∂a/∂m = 1  =  Δgm = -6
m.grad: 0 → -6
5. a → b
ga = -6  ×  ∂a/∂b = 1  =  Δgb = -6
b.grad: 0 → -6
6. m → w
gm = -6  ×  ∂m/∂w = 3  =  Δgw = -18
w.grad: 0 → -18
7. m → x
gm = -6  ×  ∂m/∂x = 2  =  Δgx = -12
x.grad: 0 → -12
value_w w value 2 grad -18 op_m × value_w->op_m ∂m/∂w=3 value_x x value 3 grad -12 value_x->op_m ∂m/∂x=2 value_m m value 6 grad -6 op_a + value_m->op_a ∂a/∂m=1 op_m->value_m value_b b value 1 grad -6 value_b->op_a ∂a/∂b=1 value_a a value 7 grad -6 op_e − value_a->op_e ∂e/∂a=1 op_a->value_a value_y y value 10 grad 6 value_y->op_e ∂e/∂y=-1 value_e e value -3 grad -6 op_L ² value_e->op_L ∂L/∂e=-6 op_e->value_e value_L L value 9 grad 1 op_L->value_L
Complete stored state after backward
The graph above stays compact. These cards expose every Value field and every saved ParentLink. The orange edge contribution is not a Value or ParentLink field; it survives only in the optional steps teaching trace.
Displayed in dependency-first order: w → x → m → b → a → y → e → L
Value w.label = 'w'
.data2.grad-18 (= ∂L/∂w).op''  —  input / leaf.parents0 saved links
.parents = ()  —  no direct operands
Value x.label = 'x'
.data3.grad-12 (= ∂L/∂x).op''  —  input / leaf.parents0 saved links
.parents = ()  —  no direct operands
Value m.label = 'm'
.data6.grad-6 (= ∂L/∂m).op'×'.parents2 saved links
.parents[0]ParentLink(value=w).valuew.local_grad3   (= ∂m/∂w)
.parents[1]ParentLink(value=x).valuex.local_grad2   (= ∂m/∂x)
Value b.label = 'b'
.data1.grad-6 (= ∂L/∂b).op''  —  input / leaf.parents0 saved links
.parents = ()  —  no direct operands
Value a.label = 'a'
.data7.grad-6 (= ∂L/∂a).op'+'.parents2 saved links
.parents[0]ParentLink(value=m).valuem.local_grad1   (= ∂a/∂m)
.parents[1]ParentLink(value=b).valueb.local_grad1   (= ∂a/∂b)
Value y.label = 'y'
.data10.grad6 (= ∂L/∂y).op''  —  input / leaf.parents0 saved links
.parents = ()  —  no direct operands
Value e.label = 'e'
.data-3.grad-6 (= ∂L/∂e).op'−'.parents2 saved links
.parents[0]ParentLink(value=a).valuea.local_grad1   (= ∂e/∂a)
.parents[1]ParentLink(value=y).valuey.local_grad-1   (= ∂e/∂y)
Value L.label = 'L'
.data9.grad1 (= ∂L/∂L).op'²'.parents1 saved link
.parents[0]ParentLink(value=e).valuee.local_grad-6   (= ∂L/∂e)

The trace contains every reverse edge. For example, the square sends \(-6\) into e.grad. On the next operation, that same stored number becomes the upstream gradient \(g_e\) for subtraction.

A single row’s product is one edge contribution to parent.grad—the quantity colored orange in our legend. If several paths return to one value, each row adds into the same buffer; only their sum is the full gradient at that parent.

Finally, check that our tiny engine and PyTorch agree at every named value.

scratch_nodes = {"w": sw, "x": sx, "m": sm, "b": sb,
                 "a": sa, "y": sy, "e": se, "L": sL}

for name, node in scratch_nodes.items():
    torch_value, torch_grad = torch_reference[name]
    assert node.data == torch_value
    assert node.grad == torch_grad

print("✓ Every value and gradient matches PyTorch.")
✓ Every value and gradient matches PyTorch.

3 · One neuron: fused sigmoid or atomic sigmoid?

Now use a slightly larger graph:

\[ m=wx,\qquad z=m+b,\qquad s=\sigma(z),\qquad e=s-y,\qquad L=e^2. \]

Choose \(w=0.5\), \(x=2\), \(b=-1\), and \(y=1\). Then \(z=0\), \(s=0.5\), and \(L=0.25\), so the backward numbers stay readable.

We will build the sigmoid in two ways:

  • fused autograd primitive: one operation \(s=\sigma(z)\);
  • atomic graph: \(n=-z\), \(q=\exp(n)\), \(d=1+q\), and \(s=1/d\).

“Fused” here describes the autograd graph: several local steps are packaged behind one operation node. It does not mean that we are skipping the chain rule.

import math

def negate(u, label):
    return Value(-u.data, label,
                 parents=(ParentLink(u, -1.0),), op="−")

def exponential(u, label):
    out = math.exp(u.data)
    return Value(out, label,
                 parents=(ParentLink(u, out),), op="exp")

def plus_one(u, label):
    return Value(1.0 + u.data, label,
                 parents=(ParentLink(u, 1.0),), op="+1")

def reciprocal(u, label):
    return Value(1.0 / u.data, label,
                 parents=(ParentLink(u, -1.0 / u.data**2),), op="1/x")

def sigmoid(u, label):
    # Stable forward formula; backward reuses the saved output s.
    if u.data >= 0:
        s = 1.0 / (1.0 + math.exp(-u.data))
    else:
        exp_z = math.exp(u.data)
        s = exp_z / (1.0 + exp_z)
    return Value(s, label,
                 parents=(ParentLink(u, s * (1.0 - s)),), op="σ")

The fused rule stores one local derivative:

\[ \frac{\partial s}{\partial z}=s(1-s). \]

The atomic graph stores four local derivatives. Their product is the same quantity:

\[ \underbrace{\left(-\frac{1}{d^2}\right)}_{s=1/d} \underbrace{(1)}_{d=1+q} \underbrace{(q)}_{q=\exp(n)} \underbrace{(-1)}_{n=-z} =\frac{q}{d^2}=s(1-s). \]

def build_sigmoid_neuron(*, fused):
    nodes = {
        "w": Value(0.5, label="w"),
        "x": Value(2.0, label="x"),
        "b": Value(-1.0, label="b"),
        "y": Value(1.0, label="y"),
    }
    nodes["m"] = multiply(nodes["w"], nodes["x"], "m")
    nodes["z"] = add(nodes["m"], nodes["b"], "z")

    if fused:
        nodes["s"] = sigmoid(nodes["z"], "s")
    else:
        nodes["n"] = negate(nodes["z"], "n")
        nodes["q"] = exponential(nodes["n"], "q")
        nodes["d"] = plus_one(nodes["q"], "d")
        nodes["s"] = reciprocal(nodes["d"], "s")

    nodes["e"] = subtract(nodes["s"], nodes["y"], "e")
    nodes["L"] = square(nodes["e"], "L")
    return nodes

fused_nodes = build_sigmoid_neuron(fused=True)
atomic_nodes = build_sigmoid_neuron(fused=False)
fused_steps = backward(fused_nodes["L"])
atomic_steps = backward(atomic_nodes["L"])

display(HTML("<h4>Fused sigmoid · 8 reverse edges</h4>"))
display(draw_graph(fused_nodes["L"], show_grad=True, min_width=1180))
show_backward_steps(fused_steps, highlight_outputs={"s"})

display(HTML("<h4 style='margin-top:24px'>Atomic sigmoid · 11 reverse edges</h4>"))
display(draw_graph(atomic_nodes["L"], show_grad=True, min_width=1700))
show_backward_steps(atomic_steps, highlight_outputs={"s", "d", "q", "n"})

Fused sigmoid · 8 reverse edges

value_w w value 0.5 grad -0.5 op_m × value_w->op_m ∂m/∂w=2 value_x x value 2 grad -0.125 value_x->op_m ∂m/∂x=0.5 value_m m value 1 grad -0.25 op_z + value_m->op_z ∂z/∂m=1 op_m->value_m value_b b value -1 grad -0.25 value_b->op_z ∂z/∂b=1 value_z z value 0 grad -0.25 op_s σ value_z->op_s ∂s/∂z=0.25 op_z->value_z value_s s value 0.5 grad -1 op_e − value_s->op_e ∂e/∂s=1 op_s->value_s value_y y value 1 grad 1 value_y->op_e ∂e/∂y=-1 value_e e value -0.5 grad -1 op_L ² value_e->op_L ∂L/∂e=-1 op_e->value_e value_L L value 0.25 grad 1 op_L->value_L
Seed: L.grad = ∂L/∂L = 1
1. L → e
gL = 1  ×  ∂L/∂e = -1  =  Δge = -1
e.grad: 0 → -1
2. e → s
ge = -1  ×  ∂e/∂s = 1  =  Δgs = -1
s.grad: 0 → -1
3. e → y
ge = -1  ×  ∂e/∂y = -1  =  Δgy = 1
y.grad: 0 → 1
4. s → z
gs = -1  ×  ∂s/∂z = 0.25  =  Δgz = -0.25
z.grad: 0 → -0.25
5. z → m
gz = -0.25  ×  ∂z/∂m = 1  =  Δgm = -0.25
m.grad: 0 → -0.25
6. z → b
gz = -0.25  ×  ∂z/∂b = 1  =  Δgb = -0.25
b.grad: 0 → -0.25
7. m → w
gm = -0.25  ×  ∂m/∂w = 2  =  Δgw = -0.5
w.grad: 0 → -0.5
8. m → x
gm = -0.25  ×  ∂m/∂x = 0.5  =  Δgx = -0.125
x.grad: 0 → -0.125

Atomic sigmoid · 11 reverse edges

value_w w value 0.5 grad -0.5 op_m × value_w->op_m ∂m/∂w=2 value_x x value 2 grad -0.125 value_x->op_m ∂m/∂x=0.5 value_m m value 1 grad -0.25 op_z + value_m->op_z ∂z/∂m=1 op_m->value_m value_b b value -1 grad -0.25 value_b->op_z ∂z/∂b=1 value_z z value 0 grad -0.25 op_n − value_z->op_n ∂n/∂z=-1 op_z->value_z value_n n value -0 grad 0.25 op_q exp value_n->op_q ∂q/∂n=1 op_n->value_n value_q q value 1 grad 0.25 op_d +1 value_q->op_d ∂d/∂q=1 op_q->value_q value_d d value 2 grad 0.25 op_s 1/x value_d->op_s ∂s/∂d=-0.25 op_d->value_d value_s s value 0.5 grad -1 op_e − value_s->op_e ∂e/∂s=1 op_s->value_s value_y y value 1 grad 1 value_y->op_e ∂e/∂y=-1 value_e e value -0.5 grad -1 op_L ² value_e->op_L ∂L/∂e=-1 op_e->value_e value_L L value 0.25 grad 1 op_L->value_L
Seed: L.grad = ∂L/∂L = 1
1. L → e
gL = 1  ×  ∂L/∂e = -1  =  Δge = -1
e.grad: 0 → -1
2. e → s
ge = -1  ×  ∂e/∂s = 1  =  Δgs = -1
s.grad: 0 → -1
3. e → y
ge = -1  ×  ∂e/∂y = -1  =  Δgy = 1
y.grad: 0 → 1
4. s → d
gs = -1  ×  ∂s/∂d = -0.25  =  Δgd = 0.25
d.grad: 0 → 0.25
5. d → q
gd = 0.25  ×  ∂d/∂q = 1  =  Δgq = 0.25
q.grad: 0 → 0.25
6. q → n
gq = 0.25  ×  ∂q/∂n = 1  =  Δgn = 0.25
n.grad: 0 → 0.25
7. n → z
gn = 0.25  ×  ∂n/∂z = -1  =  Δgz = -0.25
z.grad: 0 → -0.25
8. z → m
gz = -0.25  ×  ∂z/∂m = 1  =  Δgm = -0.25
m.grad: 0 → -0.25
9. z → b
gz = -0.25  ×  ∂z/∂b = 1  =  Δgb = -0.25
b.grad: 0 → -0.25
10. m → w
gm = -0.25  ×  ∂m/∂w = 2  =  Δgw = -0.5
w.grad: 0 → -0.5
11. m → x
gm = -0.25  ×  ∂m/∂x = 0.5  =  Δgx = -0.125
x.grad: 0 → -0.125

The blue-highlighted cards are the only part that changed:

  • fused sigmoid: one update, \(g_z=g_s\,s(1-s)=(-1)(0.25)=-0.25\);
  • atomic sigmoid: four updates, ending with the same \(g_z=-0.25\).

Fusion therefore gives a smaller graph and fewer intermediate gradient buffers. A real library can also use a numerically stable sigmoid implementation. The mathematics is unchanged: the single fused local derivative is exactly the product of the four atomic local derivatives.

common = ("w", "x", "m", "b", "z", "s", "y", "e", "L")
for name in common:
    assert math.isclose(fused_nodes[name].data, atomic_nodes[name].data)
    assert math.isclose(fused_nodes[name].grad, atomic_nodes[name].grad)

assert [(s["output"], s["parent"]) for s in fused_steps] == [
    ("L", "e"), ("e", "s"), ("e", "y"), ("s", "z"),
    ("z", "m"), ("z", "b"), ("m", "w"), ("m", "x"),
]
assert [(s["output"], s["parent"]) for s in atomic_steps] == [
    ("L", "e"), ("e", "s"), ("e", "y"), ("s", "d"),
    ("d", "q"), ("q", "n"), ("n", "z"), ("z", "m"),
    ("z", "b"), ("m", "w"), ("m", "x"),
]

expected_sigmoid = {
    "w": (0.5, -0.5), "x": (2.0, -0.125), "m": (1.0, -0.25),
    "b": (-1.0, -0.25), "z": (0.0, -0.25), "s": (0.5, -1.0),
    "y": (1.0, 1.0), "e": (-0.5, -1.0), "L": (0.25, 1.0),
}
for name, (value, grad) in expected_sigmoid.items():
    assert math.isclose(fused_nodes[name].data, value)
    assert math.isclose(fused_nodes[name].grad, grad)

fused_local = next(
    step["local"] for step in fused_steps
    if step["output"] == "s" and step["parent"] == "z"
)
atomic_locals = [
    step["local"] for step in atomic_steps
    if step["output"] in {"s", "d", "q", "n"}
]
assert math.isclose(math.prod(atomic_locals), fused_local)

tw = torch.tensor(0.5, requires_grad=True)
tx = torch.tensor(2.0, requires_grad=True)
tb = torch.tensor(-1.0, requires_grad=True)
ty = torch.tensor(1.0, requires_grad=True)
tL = (torch.sigmoid(tw * tx + tb) - ty) ** 2
tL.backward()

assert math.isclose(fused_nodes["w"].grad, tw.grad.item())
assert math.isclose(fused_nodes["x"].grad, tx.grad.item())
assert math.isclose(fused_nodes["b"].grad, tb.grad.item())
assert math.isclose(fused_nodes["y"].grad, ty.grad.item())

print("✓ Fused, atomic, and PyTorch agree.")
print("  sigmoid local: 4 atomic factors = 1 fused factor =", fused_local)
print("  final gradients: w = -0.5, x = -0.125, b = -0.25, y = 1")
✓ Fused, atomic, and PyTorch agree.
  sigmoid local: 4 atomic factors = 1 fused factor = 0.25
  final gradients: w = -0.5, x = -0.125, b = -0.25, y = 1

Takeaway

Both systems do the same three things:

  1. run the forward operations and store parent links plus local derivatives,
  2. start with \(g_L=1\) in the loss’s .grad buffer,
  3. compute upstream \(\times\) local \(=\) edge contribution to the parent, then add it to the parent’s .grad buffer.

Our tiny Value record and local rules make those steps visible. Fusion does not change the calculus; it packages a product of local derivatives behind one operation. PyTorch generalizes these ideas to tensors, neural-network layers, accelerators, and large models.