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

Feat: adjoints through observable functions #689

Merged
merged 26 commits into from
May 25, 2024
Merged

Feat: adjoints through observable functions #689

Changes from 1 commit
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
92ad6a8
feat: adjoints through observable functions
DhairyaLGandhi May 2, 2024
22dc7ec
Update ext/SciMLBaseZygoteExt.jl
DhairyaLGandhi May 6, 2024
0c2b69d
feat: allow observables in collections
DhairyaLGandhi May 8, 2024
c61e08c
chore: handle no observables in collection
DhairyaLGandhi May 8, 2024
8600a8d
fix: typo
DhairyaLGandhi May 8, 2024
a69d087
Merge branch 'master' into dg/obsfn
DhairyaLGandhi May 8, 2024
785b052
test: add test for observable functions
DhairyaLGandhi May 12, 2024
adee4f0
test: add AD testset
DhairyaLGandhi May 12, 2024
2197a30
Update test/downstream/observables_autodiff.jl
ChrisRackauckas May 13, 2024
9172014
test: add a simple DAE example; disable till sensitivities are turned on
DhairyaLGandhi May 15, 2024
95cf416
test: add missing imports
DhairyaLGandhi May 15, 2024
4ce8257
chore: format
DhairyaLGandhi May 16, 2024
839bd63
chore: rm unwanted adjoint
DhairyaLGandhi May 16, 2024
2474a8d
test: check failures with SciMLSensitivity + SII
DhairyaLGandhi May 18, 2024
f68cb05
ci(SciMLSensitivity): checkout SII branch
DhairyaLGandhi May 18, 2024
9ab29d9
ci(SciMLSensitivity): use correct path
DhairyaLGandhi May 18, 2024
a417cdd
ci: revert changes
DhairyaLGandhi May 18, 2024
ff9bb2c
chore: revert literal_getproperty adjoint
DhairyaLGandhi May 19, 2024
032b927
chore: try to avoid returning object
DhairyaLGandhi May 19, 2024
44bfc91
build: add MSL to test deps
DhairyaLGandhi May 19, 2024
de2d6cd
chore: don't return structural tangent
DhairyaLGandhi May 20, 2024
c63dfbf
chore: fix imports
DhairyaLGandhi May 20, 2024
940ea78
test: add MSL to downstream env
DhairyaLGandhi May 20, 2024
8e48f1c
test: rm MSL from test env
DhairyaLGandhi May 20, 2024
d061ce4
chore: format
DhairyaLGandhi May 20, 2024
f817b52
Update CI.yml
ChrisRackauckas May 24, 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
31 changes: 20 additions & 11 deletions ext/SciMLBaseZygoteExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
using SciMLBase: ODESolution, remake,
getobserved, build_solution, EnsembleSolution,
NonlinearSolution, AbstractTimeseriesSolution
using SymbolicIndexingInterface: symbolic_type, NotSymbolic, variable_index
using SymbolicIndexingInterface: symbolic_type, NotSymbolic, variable_index, is_observed, observed, parameter_values
using RecursiveArrayTools

# This method resolves the ambiguity with the pullback defined in
Expand Down Expand Up @@ -109,7 +109,15 @@
@adjoint function Base.getindex(VA::ODESolution, sym)
function ODESolution_getindex_pullback(Δ)
i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
if i === nothing
if is_observed(VA, sym)
y, back = Zygote.pullback(VA) do sol
f = observed(sol, sym)
p = parameter_values(sol)
f.(sol.u,Ref(p), sol.t)

Check warning on line 116 in ext/SciMLBaseZygoteExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/SciMLBaseZygoteExt.jl#L112-L116 

Added lines #L112 - L116 were not covered by tests
end
gs = back(Δ)
(gs[1], nothing)
elseif i === nothing

Check warning on line 120 in ext/SciMLBaseZygoteExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/SciMLBaseZygoteExt.jl#L118-L120 

Added lines #L118 - L120 were not covered by tests
throw(error("Zygote AD of purely-symbolic slicing for observed quantities is not yet supported. Work around this by using `A[sym,i]` to access each element sequentially in the function being differentiated."))
else
Δ′ = [[i == k ? Δ[j] : zero(x[1]) for k in 1:length(x)]
Expand All @@ -122,6 +130,7 @@

@adjoint function Base.getindex(VA::ODESolution{T}, sym::Union{Tuple, AbstractVector}) where T
function ODESolution_getindex_pullback(Δ)
@show typeof(Δ)

Check warning on line 133 in ext/SciMLBaseZygoteExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/SciMLBaseZygoteExt.jl#L133 

Added line #L133 was not covered by tests
DhairyaLGandhi marked this conversation as resolved.
Show resolved Hide resolved
sym = sym isa Tuple ? collect(sym) : sym
i = map(x -> symbolic_type(x) != NotSymbolic() ? variable_index(VA, x) : x, sym)
if i === nothing
Expand Down Expand Up @@ -182,15 +191,15 @@
NonlinearSolution{T, N, uType, R, P, A, O, uType2}(u, args...), NonlinearSolutionAdjoint
end

@adjoint function literal_getproperty(sol::AbstractTimeseriesSolution,
::Val{:u})
function solu_adjoint(Δ)
zerou = zero(sol.prob.u0)
_Δ = @. ifelse(Δ === nothing, (zerou,), Δ)
(build_solution(sol.prob, sol.alg, sol.t, _Δ),)
end
sol.u, solu_adjoint
end
# @adjoint function literal_getproperty(sol::AbstractTimeseriesSolution,
# ::Val{:u})
# function solu_adjoint(Δ)
# zerou = zero(sol.prob.u0)
# _Δ = @. ifelse(Δ === nothing, (zerou,), Δ)
# (build_solution(sol.prob, sol.alg, sol.t, _Δ),)
# end
# sol.u, solu_adjoint
# end
Copy link
Member

Choose a reason for hiding this comment

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

Why is this removed?

Copy link
Member Author

Choose a reason for hiding this comment

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

It was returning the ODESolution as the adjoint. It is also an issue because it shortcuts the gradients through parameters and instead replaces it with the sol.prob, whereas we need to accumulate the gradients here.

Copy link
Member

Choose a reason for hiding this comment

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

Can you add a unit test in the downstream set which shows this is fine?

Copy link
Member Author

Choose a reason for hiding this comment

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

Happy to. In fact, that's why I asked if anything was relying on this behavior previously. Could you suggest what kind of test you have in mind?

Copy link
Member

Choose a reason for hiding this comment

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

This seems to be the root cause of many of the test failures? So that means it's caught by the tests already.

Copy link
Member Author

Choose a reason for hiding this comment

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

I don't think this is what the error is referring to. I am missing a branch https://github.com/DhairyaLGandhi/RecursiveArrayTools.jl/tree/dg/noproj which removes an extra projection rule.

Copy link
Member Author

Choose a reason for hiding this comment

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

It does refer to projecting to a VectorOfArray, and that rule wasn't defined for Tangent. Removing it gets us the expected results. If we want to project back to a VectorOfArray type, then that needs to be handled elsewhere.

Copy link
Member Author

Choose a reason for hiding this comment

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

now that we have restored the adjoint, I believe this can be resolved


@adjoint function literal_getproperty(sol::SciMLBase.AbstractNoTimeSolution,
::Val{:u})
Expand Down
Loading