Faster ODE Adjoints with 12 Lines of Code (ICML 2021)

[adam-hartshorne]

This recent paper sets out how a simple modification to the ODE solver framework can illicit significant speed-ups.

https://github.com/patrick-kidger/FasterNeuralDiffEq

It is as simple as basically calling a different norm function on the state vector for the forward and backward pass in the initial step size and error_ratio functions.

However for the life of me, I can't work out how to code this simple adjustment to the ODE solver in Jax. It seems all that is needed is some sort of flag that says _odeint is being called in forward or backward mode and in doing so use standard norm or the alternative "norm" proposed in the paper (which basically just considers y, y_adj parameters and not the time / other arguments), but all my attempts to add this result in JAX complaining about flag parameter.