autodiff

JAX visitors for user-supplied Jacobians.

Visitors: WithJacobian

Lowers a WithJacobian node to its wrapped subexpression wrapped in a :func:jax.custom_jvp rule. The rule zeroes the tangent of every overridden variable before differentiating the inner function, then adds J @ dvar for each override — so overridden directions come from the user's Jacobian and all others from autodiff. Because the override lives inside the lowered function, every downstream differentiation (the discretizer's jacfwd, constraint linearization, sparse coloring) picks it up with no further plumbing.

Unlike the Vmap and Cond visitors, this one does not pause the value memo: the nested jax.jvp trace receives its own (x, u) tracers rather than reusing the captured outer ones, and the memo is keyed on all four arguments, so a value cached in one trace can never be handed to a caller in another.