-
-
Notifications
You must be signed in to change notification settings - Fork 123
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Updating initial guess when using nonlinear solver inside ODE term #199
Comments
Aside: this made-up example is actually quite severe, because printing Second aside: I misunderstood the documentation at https://docs.kidger.site/diffrax/api/nonlinear_solver/, I was expecting the default option |
So you'd like to pass data between vector field evaluations. It's worth noting that this isn't really a clearly-defined notion, mathematically speaking: a diffeq solver may evaluate the vector field nonmonotonically (not just forward in time), in particular when using an adaptive solver that may reject steps. That said, I agree that it can be very useful to be able to do this! Supporting this kind of side-effect hasn't been a priority for JAX so far, as side-effects are quite a complicated thing to make happen in a functional framework. Nonetheless, I think this should be possible using an upcoming JAX API, that provides for stateful operations. I've not yet tried it myself -- and it's not documented yet -- but it might suffice for this task. The operations are available here, and you can see an example of them being using in As for the nonlinear solvers, I actually have an overhaul of these planned myself. These are going to be dramatically improved soon. |
Thank you for the response - I'll keep an eye out for when that API is documented as it should be useful, but I don't currently have the time to try and figure it out for myself! Regarding variable step solvers, just for clarity, I was envisaging only updating the state after an accepted step, rather than any call to the ODE function. Along similar lines, in the case of a nonlinear solve inside the ODE function failing, it would be nice to have a mechanism to force an adaptive solver to reject that step (which is what I was doing in rtqichen/torchdiffeq#210). Good to hear about the upcoming nonlinear solver changes too! |
Hi,
I think this is partially related to #60 as it involves storing some information after an accepted step, but the difference is that I actually need to access the last known information from inside the ODE function.
I have an ODE function that requires the use of a nonlinear solver to compute the derivatives. At the moment, I'm using a fixed initial guess for
NewtonNonlinearSolver
, but this is inefficient. What I'd like to do is, after an accepted step, store the found root and use it as the initial guess during the next integration step. I was doing this in torchdiffeq successfully, but I can't see an equivalent way in Diffrax.As a (contrived) example: the code below performs some sort of nonlinear solve, but each time with a poor initial guess (meaning it takes 10 iterations to converge at each call to the ODE function). If I set
init_x = 0.9
, which is a much better guess in this case, it takes two or three iterations, so the potential benefit is clear (especially for more expensive nonlinear functions). In this case, I wouldn't expect to run into weird issues with gradients, because backpropagating throughNewtonNonlinearSolver
shouldn't depend on the initial guess.Thanks!
The text was updated successfully, but these errors were encountered: