Skip to main content

Mixed partials and coefficient budgets

A multivariate jet stores D^alpha f / alpha! for each multi-index through the requested total order. In dim dimensions and order n, the number of coefficients is binomial(dim+n, n).

Conventions​

  • mlp_jet_mv(x0, layers, order) takes one input point of shape (dim,).
  • Layer weights use (out_features, in_features) and activation=None gives an affine readout.
  • jet_partials restores ordinary partial derivatives, keyed by multi-index.
  • Coordinate ordering belongs to the caller. With (t, x), (1,0) means a time derivative and (0,2) means the second spatial derivative.

Use the mixed heat residual to see initialized tensors and parameter-gradient assertions.

Choose the representation​

Full mixed jets have combinatorial output size. Directional jets are appropriate when only a few directions are needed. A direct Laplacian path avoids forming the full Hessian or full mixed jet when only its contraction is requested.

See Torch mixed jets, JAX mixed jets, and numerical guarantees.