Owen L
Owen L
In the general case, this is basically live saving of data in a JAX while loop, I've never tried it but I assume the go to would be to have...
> 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...
Equinox modules are data classes (https://docs.python.org/3/library/dataclasses.html), so you have to specify the attributes in the class header, see https://docs.kidger.site/equinox/ for example.
Maybe this isn't even solver specific, I just noticed using Heun instead, with dt = 0.1, VBT takes 27s but UBP takes3s , so maybe the real question is twofold:...
When I compare the old to the new VBT, I didn't see a pronounced slowdown, so it doesn't seem like regression (although I am generally supportive of speed regression tests)....
> FWIW in principle JAX should be able to parallelize each call, but in practice it seems that it is not doing that... Yea, it's weird since the path is...
While it might involve some re-computation, this sounds like it could be done with a `SaveAt(fn=)`
> Unfortunately there's no good way in JAX to create something that has both custom forward-mode and custom reverse-mode autodiff. Why is this? This doesn't seem fundamentally impossible
The classic (https://github.com/patrick-kidger/diffrax/issues/446#issuecomment-2187405940) strikes once again 😉 It seems like there are a few errors here. First, you return a diagonal, but control term is for full matrices by default,...
Yes, I was using diffrax 0.6.0