PyTorch Tensors and Autograd: Leaves, VJPs and Safe Updates

The story — tensors and recorded operations

NumPy let us own the derivatives; PyTorch lets us specify tensor operations and request their pullbacks. A tensor is not literally a NumPy array: it has device, dtype, layout and autograd metadata, with interoperability rules. Eager autograd records supported operations when grad mode and input requirements call for them. It does not record every Python operation or every tensor use.

The original conference scene, framework market share and causal claims about an entire ecosystem were unsupported and are not teaching evidence. The useful historical connection is reverse-mode AD's accumulation of local derivatives, described in Baydin et al.'s survey, not a claim that every framework reimplements exactly one Wengert-list data structure.

Prerequisites, version and outcomes

Bring manual backprop, the NumPy classifier and vectorization. This lesson is checked against PyTorch2.9.0+cpu with NumPy2.2.6. No upgrade/install is required by this audit. Optional pre-reading: autograd tutorial, mechanics, Elliot Waite and Fireship.

You will create tracked leaves, inspect a dynamic graph, request scalar and vector-output gradients, port the784→128→10 network to raw tensors, update under no-grad, and diagnose missing/accumulated gradients, mutation and retained references. Compilation/FlexAttention remain source-reviewed topics, not executed speed claims.

1 · A tensor can carry a derivative history

import torch
import torch.nn.functional as F
import numpy as np
import math
 
torch.set_num_threads(1)
torch.manual_seed(14)
x=torch.tensor([1.,2.,3.],dtype=torch.float64)
assert x.shape==(3,) and (x@x).item()==14
assert torch.equal(x+1,torch.tensor([2.,3.,4.],dtype=torch.float64))
x.requires_grad_()
y=(x*x).sum();y.backward()
assert torch.equal(x.grad,2*x)

The derivative is [2,4,6] because the output sums three squares. The GPS metaphor is only an ordering aid: reverse mode applies local pullbacks, not inverse operations. A many-to-one forward operation need not be invertible to be differentiable.

A tiny bit of history: Wengert lists and reverse mode

AD propagates derivatives alongside elementary computations instead of approximating the whole program by finite differences. Wengert lists record evaluation steps; reverse-mode accumulation was developed historically before modern tensor frameworks. The survey discusses Linnainmaa and later developments. Do not infer that JAX tracing, TensorFlow graphs and PyTorch's engine have identical implementation or priority claims from this shared mathematics.

2 · requires_grad, leaves and intermediates

In default grad mode, an operation is recorded if an input requires a derivative and the operation supports it. Floating/complex differentiable tensors can require gradients; integer labels do not. A tensor with requires_grad=False is a leaf by convention. A requiring-grad tensor created directly by the user, with no tracked producer, is normally a leaf; tracked operation outputs are nonleaves.

Backward accumulates into requiring-grad leaves by default. Nonleaf gradients can be retained explicitly with retain_grad(). .grad=None can also mean unused/disconnected, no backward yet, a no-grad region, or a reset to None—not only “nonleaf.”

x=torch.tensor(3.,requires_grad=True)
y=x**2;z=y.sin();w=z+5
y.retain_grad()
w.backward()
assert abs(x.grad.item()-6*math.cos(9))<1e-5
assert abs(y.grad.item()-math.cos(9))<1e-6
assert x.is_leaf and not y.is_leaf and y.grad_fn is not None
# Scaling before enabling tracking creates a valid leaf.
good=(torch.randn(3)*.01).requires_grad_()
assert good.is_leaf
# Scaling an already tracked tensor creates a nonleaf.
base=torch.randn(3,requires_grad=True);nonleaf=base*.01
assert not nonleaf.is_leaf

2.1 A vector output needs a cotangent, not necessarily a scalar reduction

y.backward() without an explicit gradient can seed a real one-element output. It need not have shape (); a shape (1,) output is also one element. For a multi-element output pass the desired upstream vector: y.backward(v) computes vᵀJ, not the entire Jacobian. Summing or averaging is one way to choose v, not the only valid choice.

v=torch.tensor([2.,3.],dtype=torch.float64,requires_grad=True)
out=v*v
try:out.backward()
except RuntimeError:pass
else:raise AssertionError('multi-element output requires a cotangent')
out.backward(torch.tensor([4.,5.],dtype=torch.float64))
assert torch.equal(v.grad,torch.tensor([16.,30.],dtype=torch.float64))
q=torch.tensor([2.],requires_grad=True);(q*q).backward();assert q.grad.item()==4

3 · The dynamic computation graph, drawn

x=torch.tensor(2.)
y=torch.tensor(5.)
w=torch.tensor(3.,requires_grad=True)
b=torch.tensor(1.,requires_grad=True)
pred=w*x+b;err=pred-y;loss=err**2
loss.backward()
assert loss.item()==4 and w.grad.item()==8 and b.grad.item()==4

dL/derr=4, derr/dpred=1, dpred/dw=2 and dpred/db=1; therefore8 and4. Eager forward rebuilds the graph each iteration and follows the branch actually executed. It does not differentiate an unexecuted path or a Python branch condition's discontinuity.

4 · Rewriting the MNIST MLP in raw PyTorch tensors

Create the initialized values before enabling grad. (randn(...,requires_grad=True)*.01) is nonleaf; (randn(...)*.01).requires_grad_() is leaf. The old war story used the latter valid order and incorrectly called it broken.

def raw_params(d=784,h=128,k=10,seed=0):
    gen=torch.Generator().manual_seed(seed)
    W1=(torch.randn(d,h,generator=gen,dtype=torch.float64)*math.sqrt(2/d)).requires_grad_()
    W2=(torch.randn(h,k,generator=gen,dtype=torch.float64)/math.sqrt(h)).requires_grad_()
    return [W1,torch.zeros(h,dtype=torch.float64,requires_grad=True),
            W2,torch.zeros(k,dtype=torch.float64,requires_grad=True)]
 
def raw_forward(X,params):
    W1,b1,W2,b2=params
    return F.relu(X@W1+b1)@W2+b2
 
params=raw_params()
X=torch.randn(32,784,dtype=torch.float64)
y=torch.randint(0,10,(32,),dtype=torch.int64)
logits=raw_forward(X,params)
loss=F.cross_entropy(logits,y)
loss.backward()
assert logits.shape==(32,10)
assert all(p.is_leaf and p.grad.shape==p.shape for p in params)

Cross entropy expects logits and int64 class indices for this mode. It combines log-softmax and negative log likelihood with a mean reduction by default; do not feed a pre-softmaxed probability tensor as if it were logits. Soft-target mode is a separate contract. This example uses random labels/features and therefore is an API fixture, not MNIST.

5 · SGD updates and gradient reset

5.1 no_grad prevents recording the update

The in-place update of a requiring-grad leaf in grad mode normally raises, rather than silently growing a graph forever. Out-of-place differentiable optimization is another deliberate algorithm. For ordinary SGD update under torch.no_grad() after backward, then clear gradients before the next independent update.

with torch.no_grad():
    for p in params:
        p-=.1*p.grad
        p.grad=None
assert all(p.grad is None for p in params)
leaf=torch.tensor(1.,requires_grad=True)
try:leaf.add_(1)
except RuntimeError:pass
else:raise AssertionError('in-place tracked leaf must reject')

5.2 Accumulation is addition at a specified parameter point

Calling backward on fresh losses adds into existing leaf buffers. It is useful for multiple loss terms and microbatches. If parameters change between steps and old gradients are retained accidentally, their sum is not the gradient of a larger batch at the new point. Equal microbatch means need division by the number of equal-size pieces; unequal pieces need sample-count weighting. Optimizer zero_grad(set_to_none=True) differs from filling an existing buffer with zero: a None gradient can cause an optimizer to skip a parameter, including its state update.

6 · Complete raw-tensor training and local-data recipe

The following functions use neither nn.Module nor torch.optim. They can train the normalized local arrays from DL011. Keep the validation split disjoint, select by validation, then report one test score. The synthetic fixture below uses32×784 random values and a limited iteration count, explicitly not real digit training.

def raw_metrics(X,y,params,bs=512):
    if len(y)==0:raise ValueError('empty evaluation')
    total=0.;correct=0
    with torch.no_grad():
        for s in range(0,len(y),bs):
            xb,yb=X[s:s+bs],y[s:s+bs]
            z=raw_forward(xb,params)
            total+=F.cross_entropy(z,yb,reduction='sum').item()
            correct+=(z.argmax(-1)==yb).sum().item()
    return total/len(y),correct/len(y)
 
def raw_train(X,y,Xv,yv,epochs=3,bs=64,lr=.1,seed=0):
    params=raw_params(X.shape[1],128,10,seed)
    gen=torch.Generator().manual_seed(seed)
    best=None;bestloss=float('inf')
    for epoch in range(epochs):
        ids=torch.randperm(len(y),generator=gen)
        for s in range(0,len(y),bs):
            j=ids[s:s+bs]
            for p in params:p.grad=None
            loss=F.cross_entropy(raw_forward(X[j],params),y[j]);loss.backward()
            with torch.no_grad():
                for p in params:p-=lr*p.grad
        score=raw_metrics(Xv,yv,params)[0]
        if score<bestloss:
            bestloss=score;best=[p.detach().clone() for p in params]
    if best is None:raise ValueError('positive epochs required')
    return best
 
def local_mnist_raw(X_train,y_train,X_val,y_val,X_test,y_test):
    # Arrays must come from the local-only loader in DL011; no network access.
    pairs=[(X_train,y_train),(X_val,y_val),(X_test,y_test)]
    ts=[]
    for a,b in pairs:
        assert a.ndim==2 and a.shape[1]==784 and len(a)==len(b)
        ts.extend([torch.as_tensor(a,dtype=torch.float64),torch.as_tensor(b,dtype=torch.int64)])
    p=raw_train(*ts[:4])
    result=raw_metrics(*ts[4:],p)
    print('LOCAL MNIST RAW-TENSOR TEST',result)
    return p,result
 
initial=raw_metrics(X,y,params)[0]
for step in range(100):
    for p in params:p.grad=None
    loss=F.cross_entropy(raw_forward(X,params),y);loss.backward()
    with torch.no_grad():
        for p in params:p-=.1*p.grad
final=raw_metrics(X,y,params)[0]
assert final<initial and math.isfinite(final)
print('SYNTHETIC raw tensor fit only',initial,final)

The local-data function is a complete training/evaluation adapter once the explicit loader and split from DL011 have supplied the six arrays. It makes no95%-in-three-epochs promise. If matching NumPy, copy its arrays and use the same reduction/dtype; identical seed integers across libraries do not guarantee equal random tensors.

6.1 Saved tensors and graph lifetime

Autograd's engine schedules nodes when dependencies allow their pullback. Some derivatives save inputs, some outputs, some neither. Backward normally releases saved values as they are consumed; leaf .grad persists. A second backward through a saved-value-dependent graph normally raises unless it was retained. Simple graphs needing no saved tensors can sometimes be traversed again, so an error is not guaranteed for every expression.

CPU engine scheduling and GPU streams are implementation details, not a guarantee that all independent branches run simultaneously or overlap the next forward pass. A CUDA backward operation generally uses the forward operation's stream, with necessary dependencies; hardware work was not tested here.

6.2 VJP cost and dimensions

For X:(B,I), W:(I,O), Y:(B,O), the flattened Jacobian of Y with respect to W has (B*O,I*O) entries. Given G:(B,O), the VJP is X.T@G:(I,O) without constructing that Jacobian. Input VJP is G@W.T. Two matmuls explain the approximate2× backward cost for this affine map when both gradients are needed. It is not a universal bound across checkpointing, arbitrary operations or a full Jacobian request. Forward-mode JVPs propagate input directions and may be preferable for few inputs/many outputs.

7 · item, detach, clone and data

.item() returns a Python number from a one-element tensor of any shape. Printing a tensor does not itself retain it indefinitely. Appending loss tensors or summing them into another tracked tensor can retain graph structure; saved activations may already be freed after backward, so retained memory is not always the whole forward graph.

detach() cuts the derivative path but shares storage. clone() copies storage and preserves a differentiable edge when grad tracking is active. detach().clone() gives an independent, untracked snapshot. no_grad() suppresses recording for a region but does not retroactively detach existing objects; factory functions with requires_grad arguments have documented exceptions. .data can bypass mutation safety and should not be used for parameter updates.

x=torch.tensor([2.,3.],requires_grad=True)
clone=x.clone();det=x.detach();snap=x.detach().clone()
assert clone.data_ptr()!=x.data_ptr() and clone.grad_fn is not None
assert det.data_ptr()==x.data_ptr() and not det.requires_grad
with torch.no_grad():det[0]=8
assert x[0].item()==8 and snap[0].item()==2
assert torch.tensor([3.]).item()==3

For host NumPy logging use tensor.detach().cpu().numpy(); it can share CPU storage with the detached tensor, so copy if an immutable snapshot is required. CUDA .item() or host copies may synchronize; aggregate metrics on-device and transfer less often when profiling shows that cost. A detached tensor is also a valid logging value; Python numbers are not the only safe choice.

7.5 torch.compile and FlexAttention

torch.compile is opt-in, not a default mode made magically stable in a particular release. TorchDynamo captures supported Python/tensor regions, AOTAutograd can form backward graphs, and a backend such as Inductor optimizes execution. Unsupported constructs may create graph breaks; fullgraph=True requests a single graph and can error instead. Guards can trigger recompilation. Dynamic shapes and data-dependent control flow require version-specific treatment; ordinary compilation is not permission to silently accept a wrong branch.

Keep an eager reference and compare output/gradient tolerances. Measure first-call compilation separately from steady-state execution, including a representative distribution of shapes. Tiny or short runs may never amortize compilation; debug-heavy hooks/logging can change capture. No1.3–2.5× speedup is promised.

PyTorch2.9 documents FlexAttention score modification and block masks for custom attention patterns. It is a versioned prototype API; score modification and sparsity masks are distinct. Neither availability nor a code example establishes FlashAttention-2-level speed on this CPU runtime. These are source-reviewed extension points, not executed compilation/GPU tests.

8 · Failure cases — executable negative controls

Repeated backward and accumulated leaves

x=torch.tensor(2.,requires_grad=True);y=x**3
y.backward();assert x.grad.item()==12
try:y.backward()
except RuntimeError:pass
else:raise AssertionError('saved graph was consumed')
x.grad=None;y=x**3
y.backward(retain_graph=True);y.backward()
assert x.grad.item()==24

Mutation and version counters

For x*w, backward can need the forward x to differentiate w; changing x under no-grad is still detectable when that saved value is read. For x+w, the local derivative is constant, so that particular saved-value check may be unnecessary. The separate leaf-in-place rule still applies in grad mode.

x=torch.tensor(2.,requires_grad=True);w=torch.tensor(3.,requires_grad=True)
z=x*w
with torch.no_grad():x.add_(1)
try:z.backward()
except RuntimeError:pass
else:raise AssertionError('saved input mutation should be detected')
x=torch.tensor(2.,requires_grad=True);w=torch.tensor(3.,requires_grad=True)
z=x+w
with torch.no_grad():x.add_(1)
z.backward()
assert x.grad.item()==w.grad.item()==1

A version counter is efficient mutation bookkeeping, not the only imaginable detection algorithm. no_grad prevents recording the update but does not make overwriting still-needed saved values safe. Reusing a retained graph after an optimizer update is therefore hazardous.

9 · Autograd cheat sheet and decisions

• Inspect requires_grad, is_leaf, grad_fn and graph connectivity before interpreting .grad=None.
• Supply a scalar loss or an explicit cotangent for a vector output.
• Zero leaf gradients per intended optimizer step; scale accumulation by the actual reduction denominator.
• Run ordinary parameter updates under no-grad after backward.
• Use no-grad for a whole region, detach for one edge, and item/host conversion for a logging boundary. Clone when shared storage is unsafe.
• model.eval() changes module behavior, not grad mode. Some evaluation work intentionally needs derivatives.

10 · Forward builds; backward consumes saved values

The receipt metaphor is useful for provenance but incomplete: graph objects and saved tensor buffers have different lifetimes. Resetting .grad does not rebuild consumed saved tensors, and retaining a graph does not reset gradients.

11 · Retention scaffold — six answers

1. Tracking: default grad mode records supported operations involving requiring-grad inputs; nonleaf gradients need retention. Enabling grad after untracked scaling is valid and does not silently fail.
2. Hand example: loss4,dw8,db4 for x2,w3,b1,target5.
3. Failures: nonleaf/unused parameters can showNone; stale accumulators can grow or change direction; retained logging tensors can increase memory. Inspect the actual mechanism, not only symptoms.
4. VJP: dW=X.T@G, dX=G@W.T; the dense flattened Jacobian is unnecessary. The2× figure is affine-map arithmetic, not an all-program guarantee.
5. Do not compile blindly: short runs, compilation latency, unstable shapes, graph breaks and debugging can outweigh steady-state gains. Compare against eager.
6. MNIST rewrite: §§4–6 provide raw leaves, forward, CE, no-grad SGD, validation selection and a local-array adapter. The audit executed a synthetic contract, not a95% MNIST result.

Reconstruct the graph and VJP dimensions tomorrow; diagnose the negative controls next week. The next lesson packages these parameters and operations in nn.Module.

Sources and execution boundary

PyTorch2.9 autograd, cross entropy, tensor views, compiler, FlexAttention, Baydin et al.. The CPU tests validate exact lesson programs and independent gradients; no dataset download, compilation, CUDA, FlexAttention or large-model training run was performed.