Mooncake overlay for erroring out before Mooncake.DerivedRule construction #1169
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.
Closes SciML/SciMLSensitivity.jl#1230 and chalk-lab/Mooncake.jl#587.
This PR handles cases where:
DiffEqBase._concrete_solve_adjoint
must error out when usingReverseDiffAdjoint
/TrackerAdjoint
while differentiating via Mooncake. This error was already handled by SciMLSensitivity but was not getting hit.Therefore calling
Mooncake.@mooncake_overlay
forDiffEqBase.set_mooncakeoriginator_if_mooncake
is required otherwise theMooncake.DerivedRule
which contains primal typechecks fails: asTracker
for example adds tags such asTracker.TrackerReal{Float64}
aroundFloat64
's to the forward pass primals.The previous PR handled:
any other case (eg: not using
ReverseDiffAdjoint
/TrackerAdjoint
) when it is required to use aMooncake.rrule!!
forDiffEqBase.set_mooncakeoriginator_if_mooncake
in aMooncake.DerivedRule
.