diffrax
diffrax copied to clipboard
Have `ImplicitAdjoint` use `lstsq` during jvp
At the moment it uses jnp.linalg.solve
, which assumes that the solve operation is well-posed. This needn't actually be the case, though, and at the moment the gradient we obtain could be offset by anything in the kernel of the Jacobian.
(This has been fine for our other main usage of implicit differentiation -- implicit diffeq solvers -- since these are well-posed by construction.)