PyTorch: from spatial derivative to parameter gradient
Install omnibias-torch==0.5.0rc1. This example uses float64, a batch of one-dimensional inputs, a tanh hidden layer, and a linear readout.
high_order_torch.py
"""A fourth-order residual, with gradients into model parameters."""
import torch
from omnibias.torch.jet import jet_to_tower, mlp_jet
weight = torch.tensor([[0.7], [-0.4]], dtype=torch.float64, requires_grad=True)
bias = torch.tensor([0.1, 0.2], dtype=torch.float64, requires_grad=True)
readout = torch.tensor([[1.0, -0.5]], dtype=torch.float64, requires_grad=True)
layers = [(weight, bias, "tanh"), (readout, None, None)]
x = torch.linspace(-1.0, 1.0, 32, dtype=torch.float64).reshape(-1, 1)
tower = jet_to_tower(mlp_jet(x, torch.ones_like(x), layers, order=4))
assert tower.shape == (5, 32, 1)
u, du, d2u, d3u, d4u = tower.unbind(0)
loss = (d4u - 24.0).square().mean()
loss.backward()
for parameter in (weight, bias, readout):
assert parameter.grad is not None
assert torch.isfinite(parameter.grad).all()
print("tower:", tuple(tower.shape), "finite parameter gradients: True")
# This checks a residual calculation, not a solved boundary-value problem.
The derivative order occupies the first axis. Convert Taylor coefficients with jet_to_tower before treating them as ordinary derivatives. The walkthrough explains the shapes and the distinction between a residual and a solved PDE.