Automatic Differentiation: The Autograd Engine

1. How Autograd Constructs Dynamic Graphs

When a tensor is created with requires_grad=True, PyTorch begins tracking every mathematical transformation applied to it. Each operation instantiates a backward Function object (like MulBackward0 orAddBackward0) that stores the derivative formula:

import torch

# 1. Define leaf parameters with gradient tracking enabled
w = torch.tensor([2.0], requires_grad=True)
b = torch.tensor([3.0], requires_grad=True)
x = torch.tensor([4.0]) # Input data (fixed constant; does not require gradients)

# 2. Forward pass: Evaluate linear equation y = w*x + b
y = w * x + b           # y = 2.0 * 4.0 + 3.0 = 11.0

# 3. Inspect backward function nodes in the computational DAG:
print("y value:", y.item())               # 11.0
print("Grad fn of y:", y.grad_fn)         # <AddBackward0 object>
print("Next functions:", y.grad_fn.next_functions) # Links to MulBackward0 and leaf b!

# 4. Backward pass: Compute partial derivatives dy/dw and dy/db via chain rule
y.backward()

# 5. Access computed gradients:
print("dy/dw (equals x = 4.0):", w.grad.item())  # 4.0
print("dy/db (equals 1.0):    ", b.grad.item())  # 1.0

2. The Gradient Accumulation Gotcha

# Demonstration: Unintended Accumulation
x = torch.tensor([3.0], requires_grad=True)

for epoch in range(3):
    y = x ** 2
    y.backward()
    print(f"Epoch {epoch + 1} Grad (Expected 6.0): {x.grad.item()}")

# Output:
# Epoch 1: 6.0
# Epoch 2: 12.0 (Accumulated 6 + 6!)
# Epoch 3: 18.0 (Accumulated 12 + 6!)

# CORRECT FIX: Always zero gradients between steps:
x.grad.zero_()

3. Disabling Autograd for Inference: no_grad vs inference_mode

During model evaluation, validation, or live web API serving, tracking computational graphs consumes massive VRAM and slows down execution. PyTorch provides two mechanisms to disable tracking:

# 1. Detaching a tensor from the graph (returns a view sharing storage, but no grad_fn)
raw_prediction = y.detach().cpu().numpy()

# 2. Modern High-Throughput Production Inference:
with torch.inference_mode():
    features = torch.randn(64, 512)
    # Forward pass executes without any computational graph memory overhead!
    predictions = features @ torch.randn(512, 10)

4. Authoring Custom Autograd Functions

When implementing novel activation functions, specialized loss metrics, or interfacing with custom C++/CUDA kernels, subclass torch.autograd.Function and specify both forward() and backward() methods:

import torch

class CustomReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input_tensor):
        # Save input tensor for gradient calculation in the backward pass
        ctx.save_for_backward(input_tensor)
        return input_tensor.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        # Retrieve saved input tensor
        input_tensor, = ctx.saved_tensors
        # Derivative of ReLU: 1 if x > 0 else 0
        grad_input = grad_output.clone()
        grad_input[input_tensor <= 0] = 0
        return grad_input

# Apply custom operator:
relu_fn = CustomReLU.apply
z = torch.tensor([-2.0, 3.0, 5.0], requires_grad=True)
out = relu_fn(z).sum()
out.backward()

print("Input:", z)
print("Gradients:", z.grad) # tensor([0., 1., 1.])