Feat: adjoints through observable functions #689

DhairyaLGandhi · 2024-05-06T12:02:46Z

Checklist

Appropriate tests were added
Any code changes were done in a way that does not break public API
All documentation related to code changes were updated
The new code follows the
contributor guidelines, in particular the SciML Style Guide and
COLPRAC.
Any new documentation only uses public API

Additional context

Currently, ADing through observables errors, however this allows us to AD through the observable function via symbolic indexing and accumulate and return grads against sol

julia> gs3 = gradient(sol) do sol
    sum(sol[sys.w])
end
((u = [[0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0]  …  [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0]], u_analytic = nothing, errors = nothing, t = nothing, k = nothing, prob = (f = nothing, u0 = nothing, tspan = nothing, p = ([0.0, 2990.0, 0.0],), kwargs = nothing, problem_type = nothing), alg = nothing, interp = nothing, dense = nothing, tslocation = nothing, stats = nothing, alg_choice = nothing, retcode = nothing, resid = nothing, original = nothing),)

This needs handling as part of when the observable symbol is in a collection (vector/ tuple/ ...), and also for various ADs like ReverseDiff and Enzyme.

Add any other context about the problem here.

Ideally, this would be handled by removing all the adjoints related to getindex and let AD do the heavy lifting for us. But this is faster to implement in its current form.

codecov · 2024-05-06T12:08:55Z

Codecov Report

Attention: Patch coverage is 0% with 33 lines in your changes are missing coverage. Please review.

Project coverage is 29.16%. Comparing base (a0fab7a) to head (f817b52).
Report is 28 commits behind head on master.

Files	Patch %	Lines
ext/SciMLBaseZygoteExt.jl	0.00%	33 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #689      +/-   ##
==========================================
- Coverage   31.79%   29.16%   -2.64%     
==========================================
  Files          55       55              
  Lines        4535     4574      +39     
==========================================
- Hits         1442     1334     -108     
- Misses       3093     3240     +147

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

ChrisRackauckas · 2024-05-06T12:12:50Z

ext/SciMLBaseZygoteExt.jl

-@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


Why is this removed?

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.

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

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?

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

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.

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.

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

ext/SciMLBaseZygoteExt.jl

Co-authored-by: Christopher Rackauckas <[email protected]>

DhairyaLGandhi · 2024-05-08T11:26:08Z

Needs JuliaDiff/ChainRules.jl#793

ChrisRackauckas · 2024-05-08T12:32:11Z

Add your unit tests as a new downstream testset.

ChrisRackauckas · 2024-05-08T12:35:35Z

https://github.com/SciML/SciMLSensitivity.jl/blob/32f5ae7529a1957661b153f0ca9eff7e4caf0c5a/test/reversediff_output_types.jl#L14 this would hit it.

test/downstream/observables_autodiff.jl

DhairyaLGandhi · 2024-05-13T14:51:23Z

Note that with SciMLSensitivity.jl#dg/ss (and SciML/SciMLStructures.jl#18) https://github.com/SciML/SciMLSensitivity.jl/blob/32f5ae7529a1957661b153f0ca9eff7e4caf0c5a/test/reversediff_output_types.jl#L14 looks like:

julia> gs = gradient(u0 -> loss(u0), u0)
([-0.7779831009550049, 0.40028226620020263],)

DhairyaLGandhi · 2024-05-16T10:30:19Z

I've added a DAE example in the tests, but switched it off until we get SciMLSensitivity updated as well. The DC motor example fails to initialize currently. If there's a different test case, I can also hook that in.

ChrisRackauckas · 2024-05-20T02:21:03Z

Project.toml

@@ -68,6 +68,7 @@ Logging = "1.10"
 Makie = "0.20"
 Markdown = "1.10"
 ModelingToolkit = "8.75, 9"
+ModelingToolkitStandardLibrary = "2.7"


This should be gated int Downstream

I've added bounds to test/downstream/Project.toml in 940ea78, should I remove anything from the regular test environment or do i need to declare these in both places?

remove it from the regular

d061ce4 does that

DhairyaLGandhi · 2024-05-20T14:42:39Z

@ChrisRackauckas SciMLSensitivity test pass with d061ce4 (latest commit), but the Core (Downstream) tests get cancelled before anything runs. Is that because the Core (Python) tests fail for unrelated reasons?

gdalle · 2024-05-22T10:33:51Z

So what happens here is:

The latest ADTypes (v1.2) is installed in the test environment
But for the Python test group, the environment resolution forces the previous ADTypes (v0.2)
I think the following message refers to ADTypes

1 dependency precompiled but a different version is currently loaded. Restart julia to access the new version

So the latest ADTypes is provided to the Python test group even though compatibility forbids it at the moment, because the previously active environment leaks

DhairyaLGandhi · 2024-05-23T13:21:12Z

Both CI/ Python and CI/ Downgrade seem to be failing on master as well.

gdalle · 2024-05-23T13:59:58Z

The problem I mentioned has not been fixed. It's not a problem with ADTypes per se, it's a problem with environment stacking

DhairyaLGandhi · 2024-05-24T14:30:37Z

Is there anything left to be done in this PR?

feat: adjoints through observable functions

92ad6a8

ChrisRackauckas reviewed May 6, 2024

View reviewed changes

ext/SciMLBaseZygoteExt.jl Outdated Show resolved Hide resolved

DhairyaLGandhi and others added 4 commits May 6, 2024 19:33

Update ext/SciMLBaseZygoteExt.jl

22dc7ec

Co-authored-by: Christopher Rackauckas <[email protected]>

feat: allow observables in collections

0c2b69d

chore: handle no observables in collection

c61e08c

fix: typo

8600a8d

Merge branch 'master' into dg/obsfn

a69d087

DhairyaLGandhi added 2 commits May 13, 2024 00:00

test: add test for observable functions

785b052

test: add AD testset

adee4f0

ChrisRackauckas reviewed May 13, 2024

View reviewed changes

test/downstream/observables_autodiff.jl Outdated Show resolved Hide resolved

Update test/downstream/observables_autodiff.jl

2197a30

DhairyaLGandhi added 2 commits May 15, 2024 17:55

test: add a simple DAE example; disable till sensitivities are turned on

9172014

test: add missing imports

95cf416

DhairyaLGandhi closed this May 16, 2024

DhairyaLGandhi reopened this May 16, 2024

DhairyaLGandhi added 8 commits May 16, 2024 23:21

chore: format

4ce8257

chore: rm unwanted adjoint

839bd63

test: check failures with SciMLSensitivity + SII

2474a8d

ci(SciMLSensitivity): checkout SII branch

f68cb05

ci(SciMLSensitivity): use correct path

9ab29d9

ci: revert changes

a417cdd

chore: revert literal_getproperty adjoint

ff9bb2c

chore: try to avoid returning object

032b927

DhairyaLGandhi added 3 commits May 20, 2024 04:54

build: add MSL to test deps

44bfc91

chore: don't return structural tangent

de2d6cd

chore: fix imports

c63dfbf

ChrisRackauckas reviewed May 20, 2024

View reviewed changes

DhairyaLGandhi added 3 commits May 20, 2024 12:52

test: add MSL to downstream env

940ea78

test: rm MSL from test env

8e48f1c

chore: format

d061ce4

gdalle mentioned this pull request May 22, 2024

Backport #2719 to MTK v8 SciML/ModelingToolkit.jl#2726

Merged

ChrisRackauckas closed this May 23, 2024

ChrisRackauckas reopened this May 23, 2024

Update CI.yml

f817b52

ChrisRackauckas merged commit 3811745 into SciML:master May 25, 2024
29 of 42 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feat: adjoints through observable functions #689

Feat: adjoints through observable functions #689

DhairyaLGandhi commented May 6, 2024 •

edited

Loading

codecov bot commented May 6, 2024 •

edited

Loading

ChrisRackauckas May 6, 2024

DhairyaLGandhi May 6, 2024

ChrisRackauckas May 8, 2024

DhairyaLGandhi May 8, 2024

ChrisRackauckas May 16, 2024

DhairyaLGandhi May 16, 2024

DhairyaLGandhi May 16, 2024

DhairyaLGandhi May 22, 2024

DhairyaLGandhi commented May 8, 2024

ChrisRackauckas commented May 8, 2024

ChrisRackauckas commented May 8, 2024

DhairyaLGandhi commented May 13, 2024

DhairyaLGandhi commented May 16, 2024

ChrisRackauckas May 20, 2024

DhairyaLGandhi May 20, 2024

ChrisRackauckas May 20, 2024

DhairyaLGandhi May 20, 2024

DhairyaLGandhi commented May 20, 2024

gdalle commented May 22, 2024 •

edited

Loading

DhairyaLGandhi commented May 23, 2024

gdalle commented May 23, 2024

DhairyaLGandhi commented May 24, 2024

Feat: adjoints through observable functions #689

Feat: adjoints through observable functions #689

Conversation

DhairyaLGandhi commented May 6, 2024 • edited Loading

Checklist

Additional context

codecov bot commented May 6, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DhairyaLGandhi commented May 8, 2024

ChrisRackauckas commented May 8, 2024

ChrisRackauckas commented May 8, 2024

DhairyaLGandhi commented May 13, 2024

DhairyaLGandhi commented May 16, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DhairyaLGandhi commented May 20, 2024

gdalle commented May 22, 2024 • edited Loading

DhairyaLGandhi commented May 23, 2024

gdalle commented May 23, 2024

DhairyaLGandhi commented May 24, 2024

DhairyaLGandhi commented May 6, 2024 •

edited

Loading

codecov bot commented May 6, 2024 •

edited

Loading

gdalle commented May 22, 2024 •

edited

Loading