Calculating First-order Sensitivity using diffrax

Hi Patrick

I am trying to calculate the first order sensitivity w.r.t. to all the three kinetic parameters for the Robertson Equations using diffrax. I am using the stiff ODE example provided here.

As far as I understand, the only change I have to make is that I have to provide an option to calculate adjoints, so my diffeqsolve looks like the following

sol = diffrax.diffeqsolve(
        terms,
        solver,
        t0,
        t1,
        dt0,
        y0,
        saveat=saveat,
        stepsize_controller=stepsize_controller,
       adoints = diffrax.RecursiveCheckpointAdjoint(checkpoints=None)
    )

The code runs without any errors but I am not able to extract the sensitivities, and sol.ys provides values for the species, but not their sensitivities. I am trying to compare the results with the sensitivities provided by the SUNDIALS CVODES solver, see their results here.

Any help would be highly appreciated!

Start with the stiff ODE example and the diffeqsolve call, especially RecursiveCheckpointAdjoint and sol.ys. Determine how first-order sensitivities with respect to the three kinetic parameters are expected to be obtained, then verify the result against the linked CVODES Robertson-equation output.