autograd

import torch
import torch.nn as nn

# 1) Manual gradient — the core idea
x = torch.tensor(2.0, requires_grad=True)
y = x ** 3 + 4 * x          # y = 16 when x=2
y.backward()                 # dy/dx = 3x^2 + 4 = 16 when x=2
print(x.grad)                # tensor(16.)

# 2) Multiple variables
a = torch.tensor(1.0, requires_grad=True)
b = torch.tensor(2.0, requires_grad=True)
z = a ** 2 + 3 * a * b + b ** 3
z.backward()
print(a.grad)                # 2a + 3b = 8
print(b.grad)                # 3a + 3b^2 = 15

# 3) Vector gradients — backward requires a SCALAR
v = torch.randn(3, requires_grad=True)
out = v * v                # vector output
# out.backward()           # ERROR — need a scalar
out.sum().backward()         # OR pass grad of out
print(v.grad)

# 4) Disable autograd — no_grad / inference_mode
with torch.no_grad():
    pred = model(x)          # no graph built, faster, less memory

with torch.inference_mode():  # even cheaper than no_grad on PyTorch 1.9+
    pred = model(x)

# 5) Detach — break the graph
p = some_tensor.detach()     # same data, but no grad tracking
target = (input * 2).detach()  # don't backprop into target

# 6) In a model — every forward pass builds a fresh graph
model = nn.Linear(10, 1)
loss_fn = nn.MSELoss()
opt = torch.optim.SGD(model.parameters(), lr=0.01)

for x, y in batches:
    opt.zero_grad(set_to_none=True)   # fresh gradients
    pred = model(x)
    loss = loss_fn(pred, y)
    loss.backward()                    # graph walked; gradients populated
    opt.step()                         # apply gradients

# 7) zero_grad — why
# Gradients ACCUMULATE in .grad — without zeroing, every batch's gradient adds on top.
# set_to_none=True is faster + safer (None means 'no contribution').

# 8) Gradient accumulation — train with effective batch > GPU memory
opt.zero_grad(set_to_none=True)
for i, (x, y) in enumerate(batches):
    pred = model(x)
    loss = loss_fn(pred, y) / accum_steps   # normalise
    loss.backward()
    if (i + 1) % accum_steps == 0:
        opt.step()
        opt.zero_grad(set_to_none=True)

# 9) Gradient clipping — stop exploding gradients (common with RNNs)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 10) Higher-order derivatives — create_graph=True
x = torch.tensor(2.0, requires_grad=True)
y = x ** 3
grad_x = torch.autograd.grad(y, x, create_graph=True)[0]   # 3x^2 = 12
grad_grad_x = torch.autograd.grad(grad_x, x)[0]            # 6x = 12

# 11) Hooks — peek at intermediate gradients
def inspect_grad(grad):
    print('intermediate grad:', grad.norm())

z.register_hook(inspect_grad)

# Or whole-tensor hooks on Modules
def hook(module, grad_input, grad_output):
    print(module.__class__.__name__, [g.norm() if g is not None else None for g in grad_output])

for m in model.modules():
    m.register_full_backward_hook(hook)

# 12) Custom autograd Function — write your own gradient
class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.clamp(min=0)
    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        return grad_output * (x > 0).float()

y = MyReLU.apply(x)

# 13) Mixed precision — gradient scaling for fp16
scaler = torch.amp.GradScaler('cuda')
for x, y in batches:
    opt.zero_grad(set_to_none=True)
    with torch.amp.autocast('cuda', dtype=torch.float16):
        pred = model(x)
        loss = loss_fn(pred, y)
    scaler.scale(loss).backward()
    scaler.step(opt)
    scaler.update()

# 14) Common gotchas
# - Forgetting to zero gradients → loss explodes mysteriously
# - Using .item() too early breaks the graph
# - Doing `with torch.no_grad():` around the training step disables gradient calc → no learning
# - Tensor on the wrong device → silent zero gradients (move with .to(device) consistently)
# - Modifying a tensor with `requires_grad=True` in-place → autograd error or wrong gradient

# 15) Debug aids
# torch.autograd.set_detect_anomaly(True)   # slow, but pinpoints NaN/inf in backward
# torch.autograd.gradcheck(fn, (inputs,))   # verify a custom function's gradient

import torch
x = torch.tensor(3.0, requires_grad=True)
y = x ** 2 + 2 * x
y.backward()
print(x.grad)   # 2x + 2 = 8