Automatic differentiation

Automatic differentiation is the reason any of this works. The derivative $\partial_{xx} u_\theta$ is not approximated — it is computed exactly, to machine precision, by differentiating the computational graph.

Forward and reverse mode

Reverse mode costs one backward pass per output. Since a PINN typically has one scalar output and several inputs, reverse mode is the right choice, which is also what torch.autograd.grad gives you.

import torch

def residual(model, x, t, alpha=0.01):
    x.requires_grad_(True)
    t.requires_grad_(True)
    u = model(torch.cat([x, t], dim=1))

    u_t = torch.autograd.grad(u, t, torch.ones_like(u), create_graph=True)[0]
    u_x = torch.autograd.grad(u, x, torch.ones_like(u), create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x, x, torch.ones_like(u_x), create_graph=True)[0]

    return u_t - alpha * u_xx

create_graph=True on the first call is what lets you differentiate a second time. Forgetting it is the single most common reason a second-order residual comes back as None.