Jets & Faà di Bruno
A jet is a truncated Taylor expansion carried as a vector of coefficients.
omnibias propagates jets through deep compositions exactly, using the
closed-form σⁿ plus the Faà di Bruno combinatorics for composing derivatives.
The idea
If you know the Taylor jet of u at a point and you apply σ, the jet of
σ(u) is given by Faà di Bruno's formula — a sum over partitions weighted by
Bell polynomials of the inner jet, with the outer derivatives σⁿ supplied
in closed form by the tower.
(σ ∘ u)⁽ⁿ⁾ = Σ σ⁽ᵏ⁾(u) · Bₙ,ₖ(u', u'', …)
k
Because every σ⁽ᵏ⁾ is exact and every Bell polynomial is computed in
pure-Python integer/float arithmetic, the composed jet is exact to machine
precision — no nested autodiff graph.
Where the pieces live
The multi-layer jet kernels are bit-identical twins in omnibias.jax.jet
and omnibias.torch.jet. They propagate exact directional Taylor jets through
deep compositions using the closed-form σⁿ.
A minimal example
from omnibias.jax.jet import mlp_jet
# Propagate a directional Taylor jet of total order N through an MLP, getting
# every directional derivative up to order N from one structured pass.
jet = mlp_jet(x, direction=v, params=params, activation="tanh", order=4)
The directional kernel is the 1-D restriction of the multivariate jet machinery.
Why jets matter
- High-order sensitivities through deep networks without exploding autodiff graphs.
- Equation discovery: read
y, y', y'', …directly off a fitted field and search for the implicit relation (see Equation discovery). - Taylor-model enclosures: the rigorous register propagates jets with interval remainders for certified bounds (see Certified register).
Building blocks
| Kernel | Role |
|---|---|
affine_jet | jet of an affine map Wx + b |
compose_jet | Faà di Bruno composition σ ∘ u |
layer_jet | one affine + activation layer |
mlp_jet | a full multi-layer composition |
tower_to_jet / jet_to_tower | convert between derivative-tower and jet representations |
- Multivariate jets — every mixed partial from one pass.
- Faà di Bruno cookbook — a worked, certified application.