Merge pull request #221 from frankschae/ODEProblem_u0

Differentiating ArrayPartitions in DEProblems

Author: ChrisRackauckas

src/zygote.jl

@@ -19,7 +19,7 @@ function ChainRulesCore.rrule(::Type{<:ArrayPartition}, x::S, ::Type{Val{copy_x}
   function ArrayPartition_adjoint(_y)
       y = Array(_y)
       starts = vcat(0,cumsum(reduce(vcat,length.(x))))
-      NoTangent(), ArrayPartition(ntuple(i -> reshape(y[starts[i]+1:starts[i+1]], size(x[i]))), length(x)), NoTangent()
+      NoTangent(), ntuple(i -> reshape(y[starts[i]+1:starts[i+1]], size(x[i])), length(x)), NoTangent()
   end
 
   ArrayPartition(x, Val{copy_x}), ArrayPartition_adjoint

test/adjoints.jl

@@ -1,4 +1,5 @@
 using RecursiveArrayTools, Zygote, ForwardDiff, Test
+using OrdinaryDiffEq
 
 function loss(x)
     sum(abs2,Array(VectorOfArray([x .* i for i in 1:5])))
@@ -30,10 +31,17 @@ function loss5(x)
     sum(abs2,Array(ArrayPartition([x .* i for i in 1:5]...)))
 end
 
+function loss6(x)
+  _x = ArrayPartition([x .* i for i in 1:5]...)
+  _prob = ODEProblem((u,p,t)->u, _x, (0,1))
+  sum(abs2, Array(_prob.u0))
+end
+
 x = float.(6:10)
 loss(x)
 @test Zygote.gradient(loss,x)[1] == ForwardDiff.gradient(loss,x)
 @test Zygote.gradient(loss2,x)[1] == ForwardDiff.gradient(loss2,x)
 @test Zygote.gradient(loss3,x)[1] == ForwardDiff.gradient(loss3,x)
 @test Zygote.gradient(loss4,x)[1] == ForwardDiff.gradient(loss4,x)
 @test Zygote.gradient(loss5,x)[1] == ForwardDiff.gradient(loss5,x)
+@test Zygote.gradient(loss6,x)[1] == ForwardDiff.gradient(loss6,x)