Add automatic differentiation test suite and add ForwardDiff support for dynamic master solvers #455
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Krastanov
requested changes
Jun 28, 2025
Member
Krastanov
left a comment
Krastanov
left a comment
There was a problem hiding this comment.
On first pass, this looks awesome, thank you!
There are test failures though -- I have not looked into them, do you know how to address them?
Contributor
Author
|
Yes, I will do that soon, just got busy this week and didn't have time to deal with them :) |
Member
|
I will mark this as a draft, just to organize my review queue a bit. Feel free to mark it back at any time. |
Krastanov
requested changes
Nov 28, 2025
Comment on lines
+156
to
+178
| function _promote_time_and_state(u0, H::AbstractOperator, J, rates, tspan) | ||
| # TODO: Find an alternative to promote_dual, which was moved to | ||
| # an extension in DiffEqBase 6.162.0 | ||
| ext = Base.get_extension(DiffEqBase, :DiffEqBaseForwardDiffExt) | ||
| Ts = reduce(ext.promote_dual, (eltype(H), DiffEqBase.anyeltypedual(J), DiffEqBase.anyeltypedual(rates))) | ||
| Tt = real(Ts) | ||
| p = Vector{Tt}(undef,0) | ||
| u0_promote = DiffEqBase.promote_u0(u0, p, tspan[1]) | ||
| tspan_promote = DiffEqBase.promote_tspan(u0_promote.data, p, tspan, nothing, Dict{Symbol, Any}()) | ||
| return tspan_promote, u0_promote | ||
| end | ||
| _promote_time_and_state(u0, f::Function, tspan) = _promote_time_and_state(u0, f(first(tspan)..., u0), tspan) | ||
| function _promote_time_and_state(u0, f::Union{Tuple, Vector}, tspan) | ||
| # TODO: Find an alternative to promote_dual, which was moved to | ||
| # an extension in DiffEqBase 6.162.0 | ||
| ext = Base.get_extension(DiffEqBase, :DiffEqBaseForwardDiffExt) | ||
| Ts = reduce(ext.promote_dual, (eltype(f[1]), DiffEqBase.anyeltypedual.(f[2:end])...)) | ||
| Tt = real(Ts) | ||
| p = Vector{Tt}(undef,0) | ||
| u0_promote = DiffEqBase.promote_u0(u0, p, tspan[1]) | ||
| tspan_promote = DiffEqBase.promote_tspan(u0_promote.data, p, tspan, nothing, Dict{Symbol, Any}()) | ||
| return tspan_promote, u0_promote | ||
| end |
Member
There was a problem hiding this comment.
I think the piece of code linked below (that is based on something you provided) is currently the cleanest way to fix the issues stemming from promote_dual
Member
|
unmarked it from draft status to run the buildkite CI |
|
You might be interested in DifferentiationInterfaceTest.jl |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
I added a test suite that exhaustively checks automatic differentiation capabilities for each solver. With this PR, FiniteDiff.jl and ForwardDiff.jl are fully supported for schroedinger and master solvers (including their dynamic versions). I added DifferentiationInterface.jl as a test dependency to quickly test other autodiff libraries supported in Julia in the future (such as Zygote.jl and Enzyme.jl). In the future I will also add support for the stochastic, semi-classical, and Monte Carlo solvers.
One note: here I am simply testing whether or not each differentiation operation runs on each solver with random test cases. I'm open to testing for correctness within some numerical tolerance, but I'd imagine we'd have to be extremely careful that the hundreds of tests pass every time, particularly when we have a handful of autodiff libraries supported in the suite.