Skip to content

← Autograd Mechanics step 7 of 8

Medium Primitives

Count the backward graph’s nodes

Jump to the task

A tensor produced by a recorded operation exposes its backward node through grad_fn. Following that node's next_functions lets you inspect the gradient paths leading toward its inputs.

import torch

x = torch.ones(3, requires_grad=True)
y = (x * 2 + 1).sum()
print(x.is_leaf, x.grad_fn)  # True None
print(y.grad_fn is not None)  # True

For this expression, the backward path passes through sum, addition and multiplication before reaching an accumulation node for x. Node class names and graph structure are implementation details, so do not hard-code node names or assume every Python operation contributes exactly one node.

Traverse nodes, not just edges

Each next_functions entry is a pair containing another backward node and an associated index. The node may be None, for example for an input that does not require gradients. Skip those entries.

Different paths can reach the same node. For z = x * 2 followed by (z + z).sum(), both addition inputs lead to the same multiplication. Keep a visited set so each node is counted only once.

Leaves and stored gradients

A tensor requiring gradients is a leaf when it has no recorded creator operation. Tensors not requiring gradients are also leaves by convention. A leaf can participate in other tensors' graphs even though its own grad_fn is None. AccumulateGrad nodes represent accumulation into trainable leaves' grad fields.

By default, backward accumulates gradients into participating leaves. For a non-leaf intermediate, call retain_grad before backward if you want its grad field populated:

x = torch.tensor([2.], requires_grad=True)
z = x * 3
z.retain_grad()
z.square().sum().backward()
print(z.grad.tolist())  # [12.0]
print(x.grad.tolist())  # [36.0]

This is separate from retain_graph, which preserves saved values for another backward pass.

Use graph inspection as evidence

A missing path can help locate an unintended detach or no-grad boundary. Graph size alone is not a diagnosis, and a connected path does not guarantee a nonzero gradient or a parameter update. Also check gradient values and whether the optimizer holds the parameter you intend to train.

Your task

def graph_depth(y: torch.Tensor) -> int

Despite the existing function name, return a node count, not the longest path length. Count distinct backward nodes reachable from y.grad_fn, including that starting node and any reachable AccumulateGrad nodes. Skip None entries. Return zero if y.grad_fn is None, including when y itself is a trainable leaf.

The starter only inspects the first level and counts entries without checking whether they identify real, distinct nodes.