implicit function theorem

$f(\theta ,z)=0$ defines a system of nonlinear equations on $z$, parameterised by $\theta$.

Take a solution mapping function $z^{\ast }$ (e.g. solving fixed point equation with root $z^{\ast }(\theta )$), which satisfies

$f(\theta ,z^{\ast }(\theta ))=0.$

Differentiate both sides with respect to $\theta$ and evaluate at the point $({\theta }_0,z_0)$.

${\partial }_{0}f({\theta }_0,z_0)+{\partial }_{1}f({\theta }_0,z_0)\partial z^{\ast }({\theta }_0)=0$

And then rearrange for the Jacobian

$\partial z^{\ast }({\theta }_0)=-[{\partial }_{1}f({\theta }_0,z_0){]}^{-1}{\partial }_{0}f({\theta }_0,z_0).$

• The Jacobian of the solution mapping can be expressed just in terms of Jacobians of $f$ at the solution point (regardless of how we get to the solution).
• We can skip propagating derivatives through the solver using automatic differentiation systems:
  • In PyTorch, use torch.no_grad() when running the solver and add the backward pass with z.register_hook(lambda grad : torch.solve(grad[:,:,None], J.transpose(1,2))[0][:,:,0]).
  • In jax, use jax.custom_jvp and jax.custom_vjp.