Skip to main content

A fourth-order residual that trains

A high-order spatial derivative can participate in a loss while ordinary autodiff computes gradients with respect to the model parameters.

PyTorch​

Install omnibias-torch==0.5.0rc1. The layer tuples are (weight, bias, activation), with weights shaped (out_features, in_features). The final None activation is 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.

Download the complete script.

mlp_jet returns Taylor coefficients: row k is the directional derivative divided by k!. jet_to_tower restores ordinary derivatives before the residual uses the fourth row. The resulting tower has shape (5, 32, 1).

The direction is all ones because this example has one input coordinate. In multiple dimensions, a directional fourth derivative is not a Laplacian or all mixed fourth partials.

What this validates​

The assertions check the shape, parameter-gradient availability, and finite gradients. They do not claim that a boundary-value problem has been solved. A PINN workflow must also define the domain, boundary conditions, optimizer, independent validation points, and convergence criteria.

JAX​

The JAX version differentiates the residual loss with jax.value_and_grad. Enable float64 before constructing arrays when comparing against float64 references.

API contracts​

Solve a complete problem​

Continue with the executed notebook, including full initialization, training or optimization, numerical assertions, plots, and recorded results.