Multivariate jets
The single-variable jet propagates derivatives along one
direction. The multivariate generalization propagates every mixed partial
up to a total order N in a single structured forward pass.
Multi-index bookkeeping
A mixed partial in D variables is indexed by a multi-index
α = (α₁, …, α_D) with total order |α| = α₁ + … + α_D. omnibias defines the
ordering and the Cauchy-product table in pure Python:
from omnibias.core.multi_index import multi_indices, cauchy_table
idx = multi_indices(dim=3, order=2) # all α with |α| <= 2 in canonical order
The Cauchy-product table encodes how coefficients multiply when two multivariate jets are combined — the multivariate analogue of polynomial multiplication.
The kernels
Bit-identical twins live in omnibias.jax.jet_mv and
omnibias.torch.jet_mv:
| Kernel | Role |
|---|---|
identity_jet | the seed jet for the coordinate map x ↦ x |
compose_jet_mv | multivariate Faà di Bruno composition |
layer_jet_mv | one affine + activation layer, multivariate |
mlp_jet_mv | a full multi-layer composition |
jet_partials | extract a chosen set of mixed partials |
jet_gradient | extract the gradient (order-1 partials) |
jet_hessian | extract the full Hessian (order-2 partials) |
One pass, every partial
from omnibias.jax.jet_mv import mlp_jet_mv, jet_gradient, jet_hessian
jet = mlp_jet_mv(x, params=params, activation="tanh", order=2)
g = jet_gradient(jet) # ∇f(x)
H = jet_hessian(jet) # ∇²f(x), the full Hessian - not just the trace
Cost model
For total order N in dimension D, the number of multi-indices grows
combinatorially in N, but is independent of network depth per layer pass:
you pay once for the structured propagation, not once per requested partial.
This is the multivariate reason the closed-form Hessian beats a dense
autodiff Hessian — you are not re-differentiating the graph for each entry.
- Need the full Hessian or many mixed partials at a point? Use the multivariate jet.
- Need a single high-order directional derivative (or the Laplacian, which is a trace)? The directional jet or the dedicated Laplacian kernel is cheaper.