High-order derivatives in PyTorch: choose the quantity first
To compute a high-order derivative in PyTorch, first specify the quantity you need: an activation derivative, a directional derivative of a network, a collection of mixed partials, or an operator contraction such as a Laplacian. Those are different problems with different computational costs.
omnibias provides supported forward derivative paths whose outputs remain differentiable with respect to model parameters. That combination is useful when a spatial derivative must participate in a trainable physics-informed neural network (PINN) residual.
Which API matches the problem?
| Required quantity | Starting point | Runnable example |
|---|---|---|
| n-th derivative of a supported activation | Activation specification and its fastpath | First derivative |
| Selected directional derivatives of an MLP | mlp_jet and jet_to_tower | Fourth-order beam PINN |
| Mixed partial derivatives | Multivariate jet representation | Mixed-jet concepts |
| Spatial Hessian trace | deep_field_laplacian | Poisson PINN |
| Parameter-space residual curvature | A supported Gauss–Newton closure | Decay calibration |
Install the verified prerelease
python -m pip install --pre "omnibias-torch==0.5.0rc1"
This is a release candidate. Use exact versions when reproducing the examples and consult installation and compatibility for your Python and backend configuration.
Run a fourth-order derivative with parameter gradients
The following complete script builds a scalar-output tanh network, computes its first four spatial derivatives, and differentiates a fourth-order residual loss into its parameters. Save it as high_order_torch.py and run python high_order_torch.py after installing the version above.
"""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.
Expected output is tower: (5, 32, 1) finite parameter gradients: True. The leading axis contains derivative orders zero through four. This verifies a residual calculation; the beam PINN notebook adds boundary conditions, optimization, and independent solution-error checks. Download this script.
Activation derivatives are not network derivatives
An activation fast path evaluates the derivative of the activation itself. A multilayer network also needs chain-rule composition through its affine maps and intermediate activations. Forward Taylor jets provide one such representation on supported architectures.
The current layer-tuple MLP convention is (weight, bias, activation), with weights shaped (out_features, in_features). A None activation gives a linear readout. Review the Torch jet contracts before representing a different architecture.
Taylor coefficients have factorial scaling
mlp_jet(x, direction, layers, order=n) returns Taylor coefficients. The k-th coefficient is the k-th directional derivative divided by k!. Convert with jet_to_tower before inserting an ordinary derivative into a PDE formula. For a scalar-output batch, the example's converted tower has shape (n+1, batch, 1).
In one input dimension, a unit direction gives ordinary spatial derivatives. In several dimensions, a directional derivative is a contraction along that direction. It does not automatically provide every mixed partial or a repeated Laplacian.
Spatial derivatives and parameter gradients have different jobs
The spatial derivative supplies the physical operator. The parameter gradient tells the optimizer how to change the represented model. The beam notebook evaluates the fourth spatial derivative, constructs a loss, calls backward(), and asserts finite gradients into every model parameter before checking the trained solution.
Do not detach the operator result before constructing the loss. Check gradients explicitly when composing unsupported preprocessing, custom tensor operations, or external solvers.
Compute a Laplacian when you need a Laplacian
A Laplacian is a Hessian trace. Building a dense Hessian is a useful independent check on a small problem, but a supported direct contraction avoids materializing all entries when only their diagonal sum is required. The Poisson notebook uses both approaches for different jobs: a dense trace at a few verification points and a direct contraction throughout training.
Numerical and performance limits
Forward formulas remove particular nested spatial differentiation graphs; their arithmetic still has a cost. Mixed-jet coefficient counts grow combinatorially with dimension and order. Direct formulas can suffer cancellation, and larger derivative orders can amplify floating-point error.
Use float64 for the supplied reference checks. Measure other dtypes against an independent calculation on the range of inputs you actually use. Read scoped benchmark results with their batch sizes, architectures, compilation policy, and baselines. Operator timing alone does not establish end-to-end PINN training speed or PDE convergence.
A practical validation sequence
- Define the operator, units, coordinate domain, and boundary conditions.
- Verify a small derivative calculation against an independent implementation.
- Assert shapes, finite values, and gradients into intended parameters.
- Train a checkable problem with known or manufactured reference data.
- Evaluate unseen points, physical constraints, error, and runtime separately.
Explore the notebook collection · Scope and guarantees · PyTorch API