Autograd

x = torch.tensor(3.0, requires_grad=True)
y = x ** 2 + 2 * x + 1   # y = (x+1)^2
y.backward()              # dy/dx = 2x + 2 = 8 at x=3
print(x.grad)             # tensor(8.)

# Set at creation
x = torch.ones(3, requires_grad=True)
x = torch.tensor([1.0], requires_grad=True)

# Toggle after creation
x.requires_grad_(True)    # in-place
x.requires_grad_(False)

# Leaf tensors created by the user have requires_grad=False by default
# unless you set it. Intermediate (non-leaf) tensors always track grad.

print(x.is_leaf)          # True if user-created or detached
print(x.grad_fn)          # None for leaf tensors

loss = model(x)           # scalar output
loss.backward()           # computes dloss/d(all_leaves_with_requires_grad)

# Non-scalar: must supply gradient argument
t = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
out = t ** 2
out.backward(torch.ones_like(out))  # equivalent to out.sum().backward()

# Retain graph for multiple backward passes
loss.backward(retain_graph=True)    # graph not freed after first call

# Accumulate grads by calling backward() multiple times
for batch in batches:
    loss = compute_loss(batch)
    loss.backward()          # grads accumulate in .grad
optimizer.step()
optimizer.zero_grad()

optimizer.zero_grad()           # zero before each forward pass (typical)
optimizer.zero_grad(set_to_none=True)  # sets .grad = None instead of 0.0
                                       # slightly faster, safe in most cases

# Manual zero
for p in model.parameters():
    if p.grad is not None:
        p.grad.zero_()

x.grad             # accumulated gradient tensor (None until backward called)
x.grad.detach()    # gradient values without autograd tracking (never .data)
x.grad.zero_()     # manual zero

# Higher-order gradients: create_graph=True keeps grad graph alive
grad_x = torch.autograd.grad(y, x, create_graph=True)[0]
grad2_x = torch.autograd.grad(grad_x, x)[0]   # second derivative

# Disables grad tracking — no graph built, fastest for inference
with torch.no_grad():
    output = model(x)

# Decorator form
@torch.no_grad()
def predict(x):
    return model(x)

# inference_mode — stricter, even faster (views cannot be used in grad ctx)
with torch.inference_mode():
    output = model(x)

@torch.inference_mode()
def evaluate(x):
    return model(x)

y = x.detach()              # view of x's data, no grad tracking
y = x.detach().clone()      # copy, fully independent

# Common pattern: extracting a value for logging
loss_val = loss.item()      # Python float — safest
loss_cpu = loss.detach().cpu()

# Compute specific gradients explicitly
grads = torch.autograd.grad(
    outputs=loss,
    inputs=[x, y],
    grad_outputs=None,      # assumes 1.0 for scalars
    create_graph=False,
    retain_graph=False,
    allow_unused=False,     # raise if input not in graph
)
dx, dy = grads

from torch.autograd.functional import jacobian, hessian

J = jacobian(func, inputs)       # full Jacobian matrix
H = hessian(func, inputs)        # full Hessian matrix

# Efficient Jacobian-vector product (JVP)
_, jvp = torch.autograd.functional.jvp(func, inputs, v)

# Vector-Jacobian product (VJP) — what backward() does internally
_, vjp = torch.autograd.functional.vjp(func, inputs, v)

class MySigmoid(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        s = torch.sigmoid(x)
        ctx.save_for_backward(s)   # save tensors needed in backward
        return s

    @staticmethod
    def backward(ctx, grad_output):
        (s,) = ctx.saved_tensors
        return grad_output * s * (1 - s)   # chain rule

# Usage
y = MySigmoid.apply(x)

from torch.utils.checkpoint import checkpoint

# Trades compute for memory: recomputes activations during backward
out = checkpoint(layer, x)
out = checkpoint(lambda x: layer1(layer2(x)), x)

# For sequential models
out = checkpoint_sequential(layers, segments=4, input=x)

# Clip by norm (most common — prevents exploding gradients)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# Clip by value
torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)

# Check total norm without clipping
total_norm = torch.nn.utils.clip_grad_norm_(
    model.parameters(), max_norm=float('inf')
)

# Tensor hook — fires during backward, receives gradient
handle = x.register_hook(lambda grad: grad * 2)  # double the gradient
handle.remove()     # unregister

# Module forward hook
def fwd_hook(module, input, output):
    print(output.shape)
handle = layer.register_forward_hook(fwd_hook)

# Module backward hook (receives grad_input and grad_output)
def bwd_hook(module, grad_input, grad_output):
    ...
handle = layer.register_full_backward_hook(bwd_hook)
handle.remove()

# Raises informative errors for NaN/Inf in gradients — slow, use for debugging only
with torch.autograd.detect_anomaly():
    output = model(x)
    loss = criterion(output, y)
    loss.backward()

with torch.autograd.profiler.profile(use_cuda=True) as prof:
    output = model(x)
    loss = criterion(output, y)
    loss.backward()

print(prof.key_averages().table(sort_by='cpu_time_total', row_limit=10))