"""Proof: loss.backward() walks a graph that forward quietly recorded.

Builds the chapter's three-line program and prints the autograd graph
that exists before backward is ever called.
"""
import torch
import torch.nn as nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(128, 256), nn.ReLU(), nn.Linear(256, 10))

x = torch.randn(64, 128)
loss = model(x).sum()

print(f"torch {torch.__version__}")
print(f"x.grad_fn        = {x.grad_fn}")
print(f"loss.grad_fn     = {type(loss.grad_fn).__name__}")

node, depth = loss.grad_fn, 0
while node is not None and depth < 10:
    print("  " * depth + type(node).__name__)
    nexts = [n for n, _ in node.next_functions if n is not None]
    node = nexts[0] if nexts else None
    depth += 1