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)

[patrick-kidger]

I can see this is a draft, LMK when you want a review! Happy to take a look on specific points if you want direct feedback earlier.

(Although FWIW I think it'd be cleanest if #484 is the next 'big' thing we merge, just to avoid us having too many 'big' things in flight at the same time.)

[lockwo]

(Although FWIW I think it'd be cleanest if #484 is the next 'big' thing we merge, just to avoid us having too many 'big' things in flight at the same time.)

I'm good with that, although I think that's in @andyElking's control ;)

[lockwo]

One question you could answer @patrick-kidger .

Background: With this PR, the BacksolveAdjoint doesn't work with the new UBP (which is not new, it didn't before either). However, with these changes, it seems to me you it could work with static step size solvers (since you can just take the reverse of the brownian path). For adaptive solvers that reject steps, that already doesn't work for UBP, and for pre visible adaptive solvers maybe that works on the back solve (but if it takes different steps, I'm not sure it would work). My question ii, is it worth adding support for UBP with back solve adjoint (it would require some more book keeping to have the path know more about whether its being run forward/reverse, which it's basically agnostic to currently, so that it could provide the same path, but I think it would be implementable), or does this not matter as a new feature (since I have generally perceived your opinion of backsolveadjoint as detailed below)?

[patrick-kidger]

I think 'no' to adding support if it's going to be fiddly / a codepath that could potentially have more bugs. (You have correctly determined my feelings on the matter in the picture!)

But 'yes' if you think it's easy enough to do with only minor adjustments, and you feel that you'd like to add this support.

Backprop with UBP and RecursiveCheckpointAdjoint is definitely the main goal here though!

[lockwo]

Backprop with UBP and RecursiveCheckpointAdjoint is definitely the main goal here though!

Indeed, it currently works with RecursiveCheckpointAdjoint and ForwardMode (and DirectAdjoint, but not if DirectAdjoint exceeds a certain number of recursive calls, still working on that one)

[lockwo]

Ok, I'm marking this as no longer draft now. There's a few things that I still need to fix (some typing remarks I left, and the directadjoint long run failure edge case but I will need to work up the motivation to dig into bounded loop implementations this weekend lol), but I think for the most part my idea of what this PR should look like is there.

There's a few open questions, most of them are marked, but I'm putting them here too in case I deleted the comments:

• How should derivative be handled with evaluate vs call?
• Should stateful functions (as opposed to natural stateful paths) be supported in control terms?