Merge pull request #511 from SciML/enzyme_default

Fix the enzyme default check

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -23,7 +23,7 @@ jobs:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1
         with:
-          version: 1
+          version: 1.6
       - uses: actions/cache@v1
         env:
           cache-name: cache-artifacts

src/concrete_solve.jl

@@ -29,7 +29,7 @@ function DiffEqBase._concrete_solve_adjoint(prob::Union{ODEProblem,SDEProblem},
         try
           Enzyme.autodiff(Enzyme.Duplicated(du, du),
                           u0,p,prob.tspan[1]) do out,u,_p,t
-            f(out, u, _p, t)
+            prob.f(out, u, _p, t)
             nothing
           end
           true
@@ -611,7 +611,7 @@ function DiffEqBase._concrete_solve_adjoint(prob,alg,sensealg::TrackerAdjoint,
   out,pullback = Tracker.forward(tracker_adjoint_forwardpass,u0,p)
   function tracker_adjoint_backpass(ybar)
     tmp = if eltype(ybar) <: Number
-      ybar
+      Array(ybar)
     elseif typeof(ybar[1]) <: Array
       return Array(ybar)
     else

test/complex_adjoints.jl

@@ -51,3 +51,17 @@ dp2 = Zygote.gradient((p)->inner_loop(prob, p, loss_fun; sensealg = QuadratureAd
 dp3 = Zygote.gradient((p)->inner_loop(prob, p, loss_fun; sensealg = BacksolveAdjoint()), p)[1]
 @test dp1 ≈ dp2 ≈ dp3
 @test eltype(dp1) <: Float64
+
+function fiip(du,u,p,t)
+  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
+  du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
+end
+p = [1.5,1.0,3.0,1.0]; u0 = [1.0; 1.0]
+prob = ODEProblem(fiip,complex(u0),(0.0,10.0),complex(p))
+
+function sum_of_solution(u0, p)
+    _prob = remake(prob,u0=u0,p=p)
+    real(sum(solve(_prob,Tsit5(),reltol=1e-6,abstol=1e-6,saveat=0.1)))
+end
+
+dx = Zygote.gradient(sum_of_solution, complex(u0), complex(p))