Assert correctness of JVPs and VJPs of solution routines

[pnkraemer]

It may be a good idea to test the correctness of jvp(solve*) and vjp(solve*) as is currently done for solve_fixed_grid.

The reasons are threefold:

• Detect unwanted lack-of-differentiabilities in lower-level functions: e.g., are the initialisation routines both forward- and reverse-differentiable?)
• Detect NaNs
• Enable application of e.g. jax.stop_gradient without destroying functionality.

This is especially important because one current application of this code seems to be parameter estimation, which heavily relies on gradients.