The node behind grad_fn

The chapter above this lesson says y.grad_fn is the tape, which is true and stops one step short of useful. type(y.grad_fn) comes back as SumBackward0, a class generated at build time from the derivative formula table, and the object is an instance of the C++ torch::autograd::Node base exposed to Python. It answers name(), it carries a next_functions tuple, and for most ops it carries the tensors it saved under _saved_ attributes that lesson three takes apart.

Two names for the same node disagree on purpose. The repr prints the short class name, AccumulateGrad, while name() returns the fully qualified torch::autograd::AccumulateGrad, because the leaf accumulator is a hand-written C++ node rather than a generated one. When you are matching a profiler trace against a graph walk, the long name is the one the profiler emitted.

Run the walk on the simplest expression that branches and something shows up that a drawing would get wrong. x * x has two operands, so MulBackward0 holds two edges, and both edges point at the same AccumulateGrad object. Not two accumulators for one tensor. One accumulator, reached twice, whose .variable attribute is x itself.

import torch

x = torch.ones(3, requires_grad=True)
y = (x * x).sum()
mul = y.grad_fn.next_functions[0][0]
print(y.grad_fn.name())              # SumBackward0
print(len(mul.next_functions))       # 2
a, b = (n for n, _ in mul.next_functions)
print(a.name())                      # torch::autograd::AccumulateGrad
print(a is b, a.variable is x)       # True True

next_functions is a tuple of pairs, and both halves matter. The first half is the node this edge leads to, or None when the operand on that side has no gradient to receive. The second half is an integer nobody looks at until they need it, and the table below is the cheapest way to learn what fills each position.

Read the None rows first. Multiplying by a constant produces the same MulBackward0 in both spellings, but the edge for the constant side is (None, 0), so the engine has somewhere to put the operand's slot without having anywhere to send a gradient. That is how autograd represents a dead branch: not by omitting the edge, but by keeping the position and nulling the destination.

The chapter's tape exhibit walks four of these chains node by node, and two of its four entries are the pair worth holding side by side. The one-leaf multiply and the two-leaf matmul both print a node with two AccumulateGrad children, and they are not the same shape underneath. In the multiply, both children are one object; in the matmul, they are two objects with different .variable tensors.

Every row in that table has a zero in the slot position, which is what happens when every op you try produces one output. Split a tensor and the zero stops being automatic. x.split(3) returns two tensors that share one SplitBackward0 node, and they are told apart by output_nr: the first is output zero, the second is output one.

Now the pair makes sense. An edge says which node to call and which of that node's outputs this gradient belongs to, because a node with two outputs receives two gradients and has to know which is which. Take the gradient of something built from the second half of the split and the edge reads ('SplitBackward0', 1).

One tensor attribute carries the same number on the forward side. hi.output_nr is 1 before any backward pass runs, so you can read a tensor's position in its producing node without touching the graph at all. It is also the number that appears in the version-counter error the museum wing keeps, in the phrase naming which output of which node was modified.

import torch

x = torch.arange(6.0, requires_grad=True)
lo, hi = x.split(3)
print(lo.grad_fn is hi.grad_fn)      # True: one node, two outputs
print(lo.output_nr, hi.output_nr)    # 0 1
edge = (hi * 2).sum().grad_fn.next_functions[0][0].next_functions
print([(n.name() if n else None, k) for n, k in edge])
# [('SplitBackward0', 1), (None, 0)]

Each node gets an integer when it is constructed, from a counter that starts at zero in every thread and only goes up. Four ops in forward order come out zero, one, two, three, which reads like bookkeeping until you notice which node breaks the pattern.

AccumulateGrad is not on the counter. Its number is the largest a 64-bit unsigned integer holds, 18446744073709551615, set by hand rather than drawn. A leaf accumulator is created whenever a leaf is first used, so on the counter it would land wherever that happened to be, and its position on the counter is exactly what the engine sorts by.

You can read the number off any node with _sequence_nr(). The underscore is honest about the stability promise, and the value is worth reading anyway, because it is the same integer the profiler stamps on a backward event when it correlates that event with the forward op that produced it.

import torch

x = torch.ones(3, requires_grad=True)
a, b = x * 2, x * 3
c = a + b
y = c.sum()
for t in (a, b, c, y):
    print(t.grad_fn.name(), t.grad_fn._sequence_nr())
# MulBackward0 0 / MulBackward0 1 / AddBackward0 2 / SumBackward0 3
acc = c.grad_fn.next_functions[0][0].next_functions[0][0]
print(acc.name(), acc._sequence_nr())
# torch::autograd::AccumulateGrad 18446744073709551615

The comment above sequence_nr() in the C++ header states both jobs the number does, and the first one turns the measurement above into a scheduling rule. Higher runs first, so the op that ran last in the forward pass is the first one the backward pass reaches. Reverse order falls out of a priority queue rather than out of a stored ordering.

The caveat in the middle of that comment is the reason AccumulateGrad sits at the ceiling. Give the leaf accumulator the maximum priority and a gradient that has arrived at a leaf gets written the moment it can be, instead of waiting behind interior nodes that were created later. The buffer holding that gradient is released earlier as a result.

That is one number doing scheduling and profiling at once, and it is worth knowing that the second job is why it exists at all in a profiled build. The rest of this arc stays on the first job: which node runs, what it reads, and what it writes.

  /// NOTE [ Sequence Number]
  ///
  /// The sequence_nr has two main usages in autograd:
  ///
  /// 1) Helps determine the node's execution priority in the engine.
  ///    All else being equal, nodes with higher priority numbers are executed
  ///    first. Thus, nodes corresponding to ops executed later are the first to
  ///    be executed in the backward pass. One caveat is that we prioritize
  ///    AccumulateGrad nodes by explicitly setting its sequence_nr to be
  ///    UINT64_MAX.