Skip to content
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.

Sign up for GitHub

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

Allow BackSolveAdjoint to accept MTKParameters #1131

Draft
wants to merge 3 commits into
base: master
Choose a base branch
from
Draft

Allow BackSolveAdjoint to accept MTKParameters #1131

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b161c2b
chore: allow mtkp in BacksolveADjoint
DhairyaLGandhi Oct 7, 2024
e8ef956
chore: rm Parameter check for sensealg
DhairyaLGandhi Oct 7, 2024
bb5be4e
Merge branch 'master' into dg/bs_adjoint
DhairyaLGandhi Nov 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion src/concrete_solve.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -634,9 +634,19 @@ function DiffEqBase._concrete_solve_adjoint(

du0 = reshape(du0, size(u0))

dp = p === nothing || p === SciMLBase.NullParameters() ? nothing :
dp = p === nothing || p === DiffEqBase.NullParameters() ? nothing :
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It should be SciMLBase

Copy link
Member Author

@DhairyaLGandhi DhairyaLGandhi Nov 11, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we can move the refactoring in its own PR. It only happened because the last commit was for getting rid of some conflicts and the merge commit had that diff. Its not in the diff for this PR itself. We need to get some tests for the adjoint paths we expose in this PR.

dp isa AbstractArray ? reshape(dp', size(tunables)) : dp

_, repack_adjoint = if p === nothing || p === DiffEqBase.NullParameters() ||
!isscimlstructure(p)
nothing, x -> (x,)
else
Zygote.pullback(p) do p
t, _, _ = canonicalize(Tunable(), p)
t
end
end

if originator isa SciMLBase.TrackerOriginator ||
originator isa SciMLBase.ReverseDiffOriginator
(NoTangent(), NoTangent(), du0, repack_adjoint(dp)[1], NoTangent(),
Expand Down
Loading
You are viewing a condensed version of this merge commit. You can view the full changes here.