Stateful Controls

[lockwo]

This is the v1 of the stateful control PR.

This PR introduces a suite of internally breaking changes which allow for stateful paths. To the non-advanced user, this will likely have no impact on the way they interact with diffrax. It should not impact ODE users at all. SDEs users may be impacted as they might want to try the new "UnsafeBrownianPath" (or as I am calling it "DirectBrownianPath" since it isn't really that unsafe any more), but this should just allow for faster code. I tried several approaches to this problem, but I think based on the conversations in other issues (e.g. #548), that this is at least someone an aligned vision.

Functionaly, this was just a lot of adding arguments to things to support an extra path_state variable. Hence the internal breaking changes. However, these breaking changes come with some decent motivation:

• Flexibility and speed: this allows for SDE users to see meaningful speedups (on smaller problems where the VF isn't something massive). VBT is necessarily (and unecessarily VBT vs brownian path slowdown #489) slower than UBP, and in theory is only necessary for adaptive solvers that reject steps (I think? Also maybe stateful paths allow some computations to carry brownian bridge states to work with rejecting steps, I think Frank had some jax code that did that but wasn't in diffrax). Thus for anyone running non adaptive solvers (or running solvers that use pre-visible step size adaptation, not currently in diffrax but we have several of those), VBT is not necessary. In addition to the speedups users get from being able to use UBP where before they had to use VBT, UBP is also faster. UBP is faster in both of its new modes, one in which it just splits the key as it goes along and (much faster) in the case where we precompute brownian motion. This addresses Performance issue with SDE solver #517 and Why can't UBP be backproped? #490.
• Generalization: paths can be stateful, and this is a more complete and generalizable picture that will likely unlock new functionality (and previously, the UBP relied on a bit of a hack to split the keys to not be stateful). Being applied at the AbstractPath level also checks off one point of Tidy-up to-dos #217.
• Also fixes Residual typing issues in ULD #570

Currently, this PR is missing a few things though.

• Direct adjoint edge case bug
• There's some comments/questions I left for myself (and reviewers) that I still need to address (I can extract some of the main ones to be sure to get feedback)

For more context, see the original PR: #559

[alessandrofasse]

I would be very interested in this. Is there any plan to continue on it :)?

[lockwo]

I would be very interested in this. Is there any plan to continue on it :)?

I will do my best to keep it up to date (as I do with all my PRs, even those that are a year old #463). We have some internal tools that also depend on this branch (so I think it should be a little stable).

I think the only outstanding real bug/issue is the weird edge case that cause a assertion error with direct adjoint on problems with enough steps (which I haven't dug into since I haven't been working a ton with diffrax this year so far).

However, this introduces quite some changes, so the timeline is long.

I have a fork of diffrax that has additional features/solvers and I was thinking this could be part of it as well, although with how much other stuff is going on, I'm not sure that timeline is any shorter lol.

[alessandrofasse]

Dear Lockwo,

thanks for the comment and the explanation. I can imagine that this has a long tail. I think I am not enough into the topic to support you here. But in any case, you work is much appreciated! If there are any updates I would be happy to use it. As of now I think I need to drop diffrax for stochastic integration because it introduces for me a factor of being 20x slower, which is in my case a lot.

[patrick-kidger]

Yup, this branch is still something I have an eye on. I'm tentatively thinking that after reversible adjoints are in (#603) then this is a good candidate for the next 'big' thing I'm thinking about.

Right now I'm not sure if that will mean stateful controls or something else, since the main downstream goal is a performance improvement for SDEs, and perhaps this is a thing we can obtain in some other way. Either way, a problem for future Patrick 😄