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_adjointmust error out when usingReverseDiffAdjoint/TrackerAdjointwhile differentiating via Mooncake. This error was already handled by SciMLSensitivity but was not getting hit.Therefore calling
Mooncake.@mooncake_overlayforDiffEqBase.set_mooncakeoriginator_if_mooncakeis required otherwise theMooncake.DerivedRulewhich contains primal typechecks fails: asTrackerfor 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_mooncakein aMooncake.DerivedRule.