Partial implementation of Zygote AD with symbolic indexing

Project.toml

@@ -6,6 +6,7 @@ version = "1.94.0"
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
@@ -28,6 +29,13 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
+
+[weakdeps]
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[extensions]
+ZygoteExt = "Zygote"
 
 [compat]
 ADTypes = "0.1.3"
@@ -51,6 +59,8 @@ SymbolicIndexingInterface = "0.2"
 Tables = "1"
 TruncatedStacktraces = "1"
 julia = "1.6"
+ChainRulesCore = "1.16"
+ZygoteRules = "0.2"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
@@ -61,6 +71,7 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Pkg", "SafeTestsets", "Test", "StaticArrays", "StochasticDiffEq", "Aqua"]
+test = ["Pkg", "SafeTestsets", "Test", "StaticArrays", "StochasticDiffEq", "Aqua", "Zygote"]

ext/ZygoteExt.jl

@@ -0,0 +1,58 @@
+module ZygoteExt
+
+using Zygote: pullback
+using ZygoteRules: @adjoint
+using SciMLBase: ODESolution, issymbollike, sym_to_index, remake, getobserved
+
+# This method resolves the ambiguity with the pullback defined in
+# RecursiveArrayToolsZygoteExt
+# https://github.com/SciML/RecursiveArrayTools.jl/blob/d06ecb856f43bc5e37cbaf50e5f63c578bf3f1bd/ext/RecursiveArrayToolsZygoteExt.jl#L67
+@adjoint function getindex(VA::ODESolution, i::Int, j::Int)
+    function ODESolution_getindex_pullback(Δ)
+        du = [m == j ? [i == k ? Δ : zero(VA.u[1][1]) for k = 1:length(VA.u[1])] :
+              zero(VA.u[1]) for m = 1:length(VA.u)]
+        dp = zero(VA.prob.p)
+        dprob = remake(VA.prob, p = dp)
+        du, dprob
+        T = eltype(eltype(VA.u))
+        N = length(VA.prob.p)
+        Δ′ = ODESolution{T,N,typeof(du),Nothing,Nothing,typeof(VA.t),
+            typeof(VA.k),typeof(dprob),typeof(VA.alg),typeof(VA.interp),
+            typeof(VA.destats),typeof(VA.alg_choice)}(du, nothing, nothing,
+            VA.t, VA.k, dprob, VA.alg, VA.interp, VA.dense, 0, VA.destats,
+            VA.alg_choice, VA.retcode)
+        (Δ′, nothing, nothing)
+    end
+    VA[i, j], ODESolution_getindex_pullback
+end
+
+@adjoint function getindex(VA::ODESolution, sym, j::Int)
+    function ODESolution_getindex_pullback(Δ)
+        i = issymbollike(sym) ? sym_to_index(sym, VA) : sym
+        du, dprob = if i === nothing
+            getter = getobserved(VA)
+            grz = pullback(getter, sym, VA.u[j], VA.prob.p, VA.t[j])[2](Δ)
+            du = [k == j ? grz[2] : zero(VA.u[1]) for k = 1:length(VA.u)]
+            dp = grz[3] # pullback for p
+            dprob = remake(VA.prob, p = dp)
+            du, dprob
+        else
+            du = [m == j ? [i == k ? Δ : zero(VA.u[1][1]) for k = 1:length(VA.u[1])] :
+                  zero(VA.u[1]) for m = 1:length(VA.u)]
+            dp = zero(VA.prob.p)
+            dprob = remake(VA.prob, p = dp)
+            du, dprob
+        end
+        T = eltype(eltype(VA.u))
+        N = length(VA.prob.p)
+        Δ′ = ODESolution{T,N,typeof(du),Nothing,Nothing,typeof(VA.t),
+            typeof(VA.k),typeof(dprob),typeof(VA.alg),typeof(VA.interp),
+            typeof(VA.destats),typeof(VA.alg_choice)}(du, nothing, nothing,
+            VA.t, VA.k, dprob, VA.alg, VA.interp, VA.dense, 0, VA.destats,
+            VA.alg_choice, VA.retcode)
+        (Δ′, nothing, nothing)
+    end
+    VA[sym, j], ODESolution_getindex_pullback
+end
+
+end

src/SciMLBase.jl

@@ -21,6 +21,8 @@ import RuntimeGeneratedFunctions
 import EnumX
 import TruncatedStacktraces
 import ADTypes: AbstractADType
+import ChainRulesCore
+import ZygoteRules: @adjoint
 
 using Reexport
 using SciMLOperators
@@ -701,6 +703,7 @@ include("solutions/optimization_solutions.jl")
 include("solutions/dae_solutions.jl")
 include("solutions/pde_solutions.jl")
 include("solutions/solution_interface.jl")
+include("solutions/zygote.jl")
 
 include("ensemble/ensemble_solutions.jl")
 include("ensemble/ensemble_problems.jl")

src/solutions/chainrules.jl

@@ -0,0 +1,72 @@
+function ChainRulesCore.rrule(config::ChainRulesCore.RuleConfig{
+        >:ChainRulesCore.HasReverseMode,
+    },
+    ::typeof(getindex),
+    VA::ODESolution,
+    sym,
+    j::Integer)
+    function ODESolution_getindex_pullback(Δ)
+        i = issymbollike(sym) ? sym_to_index(sym, VA) : sym
+        if i === nothing
+            getter = getobserved(VA)
+            grz = rrule_via_ad(config, getter, sym, VA.u[j], VA.prob.p, VA.t[j])[2](Δ)
+            du = [k == j ? grz[2] : zero(VA.u[1]) for k in 1:length(VA.u)]
+            dp = grz[3] # pullback for p
+            dprob = remake(VA.prob, p = dp)
+            T = eltype(eltype(VA.u))
+            N = length(VA.prob.p)
+            Δ′ = ODESolution{T, N, typeof(du), Nothing, Nothing, Nothing, Nothing,
+                typeof(dprob), Nothing, Nothing, Nothing, Nothing}(du, nothing,
+                nothing, nothing, nothing, dprob, nothing, nothing,
+                VA.dense, 0, nothing, nothing, VA.retcode)
+            (NoTangent(), Δ′, NoTangent(), NoTangent())
+        else
+            du = [m == j ? [i == k ? Δ : zero(VA.u[1][1]) for k in 1:length(VA.u[1])] :
+                  zero(VA.u[1]) for m in 1:length(VA.u)]
+            dp = zero(VA.prob.p)
+            dprob = remake(VA.prob, p = dp)
+            Δ′ = ODESolution{
+                T,
+                N,
+                typeof(du),
+                Nothing,
+                Nothing,
+                typeof(VA.t),
+                typeof(VA.k),
+                typeof(dprob),
+                typeof(VA.alg),
+                typeof(VA.interp),
+                typeof(VA.alg_choice),
+                typeof(VA.destats),
+            }(du,
+                nothing,
+                nothing,
+                VA.t,
+                VA.k,
+                dprob,
+                VA.alg,
+                VA.interp,
+                VA.dense,
+                0,
+                VA.destats,
+                VA.alg_choice,
+                VA.retcode)
+            (NoTangent(), Δ′, NoTangent(), NoTangent())
+        end
+    end
+    VA[sym, j], ODESolution_getindex_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(getindex), VA::ODESolution, sym)
+    function ODESolution_getindex_pullback(Δ)
+        i = issymbollike(sym) ? sym_to_index(sym, VA) : sym
+        if i === nothing
+            throw("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)]
+                  for (x, j) in zip(VA.u, 1:length(VA))]
+            (NoTangent(), Δ′, NoTangent())
+        end
+    end
+    VA[sym], ODESolution_getindex_pullback
+end

src/solutions/zygote.jl

@@ -0,0 +1,22 @@
+@adjoint function getindex(VA::ODESolution, i::Int)
+    function ODESolution_getindex_pullback(Δ)
+        Δ′ = [[i == k ? Δ[j] : zero(x[1]) for k in 1:length(x)]
+              for (x, j) in zip(VA.u, 1:length(VA))]
+        (Δ′, nothing)
+    end
+    VA[i], ODESolution_getindex_pullback
+end
+
+@adjoint function getindex(VA::ODESolution, sym)
+    function ODESolution_getindex_pullback(Δ)
+        i = issymbollike(sym) ? sym_to_index(sym, VA) : sym
+        if i === nothing
+            throw("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)]
+                  for (x, j) in zip(VA.u, 1:length(VA))]
+            (Δ′, nothing)
+        end
+    end
+    VA[sym], ODESolution_getindex_pullback
+end

test/downstream/Project.toml

@@ -4,6 +4,7 @@ ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"

test/downstream/remake_autodiff.jl

@@ -1,27 +1,47 @@
 using OrdinaryDiffEq, ModelingToolkit, Zygote, SciMLSensitivity
 
-@variables t
+@variables t x(t) o(t)
 D = Differential(t)
 function lotka_volterra(; name = name)
-    states = @variables x(t)=1.0 y(t)=1.0
+    states = @variables x(t)=1.0 y(t)=1.0 o(t)
     params = @parameters p1=1.5 p2=1.0 p3=3.0 p4=1.0
     eqs = [
         D(x) ~ p1 * x - p2 * x * y,
         D(y) ~ -p3 * y + p4 * x * y,
+        o ~ x * y,
     ]
     return ODESystem(eqs, t, states, params; name = name)
 end
 
 @named lotka_volterra_sys = lotka_volterra()
+lotka_volterra_sys = structural_simplify(lotka_volterra_sys)
 prob = ODEProblem(lotka_volterra_sys, [], (0.0, 10.0), [])
 sol = solve(prob, Tsit5(), reltol = 1e-6, abstol = 1e-6)
+u0 = [1.0 1.0]
+p = [1.5 1.0 1.0 1.0]
 
 function sum_of_solution(u0, p)
     _prob = remake(prob, u0 = u0, p = p)
     sum(solve(_prob, Tsit5(), reltol = 1e-6, abstol = 1e-6, saveat = 0.1,
         sensealg = BacksolveAdjoint(autojacvec = ZygoteVJP())))
 end
 
-u0 = [1.0 1.0]
-p = [1.5 1.0 1.0 1.0]
 du01, dp1 = Zygote.gradient(sum_of_solution, u0, p)
+
+function symbolic_indexing(u0, p)
+    _prob = remake(prob, u0 = u0, p = p)
+    soln = solve(_prob, Tsit5(), reltol = 1e-6, abstol = 1e-6, saveat = 0.1,
+        sensealg = BacksolveAdjoint(autojacvec = ZygoteVJP()))
+    sum(soln[x])
+end
+
+du01, dp1 = Zygote.gradient(symbolic_indexing, u0, p)
+
+function symbolic_indexing_observed(u0, p)
+    _prob = remake(prob, u0 = u0, p = p)
+    soln = solve(_prob, Tsit5(), reltol = 1e-6, abstol = 1e-6, saveat = 0.1,
+        sensealg = BacksolveAdjoint(autojacvec = ZygoteVJP()))
+    sum(soln[o, i] for i in 1:length(soln))
+end
+
+du01, dp1 = Zygote.gradient(symbolic_indexing_observed, u0, p)