Experiments to try different AD systems

[SamuelBrand1]

AD system with ODEs

There are two main ways to differentiate functions that depend on the solution to an ODE initial value problem where the ODE solution itself depends on some parameters $p$:

• Create a bigger ODE system where the sensitivity of the ODE solution to $p$ at any time $t$ can be calculated by solving the larger system. This can be forward sensitivity or adjoint sensitivity (cf https://docs.sciml.ai/SciMLSensitivity/stable/sensitivity_math/).
• Treat the ODE solver as itself a complex sequence of primitive operations and AD through it directly, this is sometimes called "Discretise then optimise" e.g https://docs.kidger.site/diffrax/api/adjoints/ .

Experiments

Apols if I've missed it, but I think it would be a good idea to experiment with the different adjoint options in diffrax in combination with the grad options in numpyro. In my experience this can sometimes change reliability and performance.