chore: revert bad commit

Author: DhairyaLGandhi

src/adjoint_common.jl

@@ -749,16 +749,16 @@ function out_and_ts(_ts, duplicate_iterator_times, sol)
     return out, ts
 end
 
-# if !hasmethod(Zygote.adjoint,
-#     Tuple{Zygote.AContext, typeof(Zygote.literal_getproperty),
-#         SciMLBase.AbstractTimeseriesSolution, Val{:u}})
-#     Zygote.@adjoint function Zygote.literal_getproperty(sol::AbstractTimeseriesSolution,
-#             ::Val{:u})
-#         function solu_adjoint(Δ)
-#             zerou = zero(sol.prob.u0)
-#             _Δ = @. ifelse(Δ === nothing, (zerou,), Δ)
-#             (SciMLBase.build_solution(sol.prob, sol.alg, sol.t, _Δ),)
-#         end
-#         sol.u, solu_adjoint
-#     end
-# end
+if !hasmethod(Zygote.adjoint,
+    Tuple{Zygote.AContext, typeof(Zygote.literal_getproperty),
+        SciMLBase.AbstractTimeseriesSolution, Val{:u}})
+    Zygote.@adjoint function Zygote.literal_getproperty(sol::AbstractTimeseriesSolution,
+            ::Val{:u})
+        function solu_adjoint(Δ)
+            zerou = zero(sol.prob.u0)
+            _Δ = @. ifelse(Δ === nothing, (zerou,), Δ)
+            (SciMLBase.build_solution(sol.prob, sol.alg, sol.t, _Δ),)
+        end
+        sol.u, solu_adjoint
+    end
+end

src/gauss_adjoint.jl

@@ -488,19 +488,8 @@ function vec_pjac!(out, λ, y, t, S::GaussIntegrand)
         ReverseDiff.reverse_pass!(tape)
         copyto!(vec(out), ReverseDiff.deriv(tp))
     elseif sensealg.autojacvec isa ZygoteVJP
-        # global gf = f
-        # global gy = y
-        # global gtunables = tunables
-        # global grepack = repack
-        # global gt = t
-        # @show f(y, tunables, t)
-        c = Zygote.bufferfrom(y)
-        c = copy(y)
-        # @show f(c, y, repack(tunables), t)
         _dy, back = Zygote.pullback(tunables) do tunables
-            c = Zygote.bufferfrom(y)
-            f(c, y, repack(tunables), t)
-            vec(copy(c))
+            vec(f(y, repack(tunables), t))
         end
         tmp = back(λ)
         if tmp[1] === nothing

src/quadrature_adjoint.jl

@@ -176,7 +176,6 @@ function AdjointSensitivityIntegrand(sol, adj_sol, sensealg, dgdp = nothing)
     adj_prob = adj_sol.prob
     (; f, tspan) = prob
     p = parameter_values(prob)
-    tunables, _, _ = canonicalize(Tunable(), p)
     u0 = state_values(prob)
     numparams = length(p)
     y = zero(state_values(prob))
@@ -253,9 +252,6 @@ function AdjointSensitivityIntegrand(sol, adj_sol, sensealg, dgdp = nothing)
     end
     AdjointSensitivityIntegrand(sol, adj_sol, p, y, λ, pf, f_cache, pJ, paramjac_config,
         sensealg, dgdp_cache, dgdp)
-
-    # AdjointSensitivityIntegrand(sol, adj_sol, tunables, y, λ, pf, f_cache, pJ, paramjac_config,
-    #     sensealg, dgdp_cache, dgdp)
 end
 
 # out = λ df(u, p, t)/dp at u=y, p=p, t=t
@@ -289,15 +285,13 @@ function vec_pjac!(out, λ, y, t, S::AdjointSensitivityIntegrand)
         copyto!(vec(out), ReverseDiff.deriv(tp))
     elseif sensealg.autojacvec isa ZygoteVJP
         _dy, back = Zygote.pullback(p) do p
-            # @show f(y, p, t)
             vec(f(y, p, t))
         end
         tmp = back(λ)
-        # @show tmp
         if tmp[1] === nothing
             out[:] .= 0
         else
-            out[:] .= vec(tmp[1].tunable)
+            out[:] .= vec(tmp[1])
         end
     elseif sensealg.autojacvec isa MooncakeVJP
         _, _, p_grad = mooncake_run_ad(paramjac_config, y, p, t, λ)
@@ -336,7 +330,6 @@ function (S::AdjointSensitivityIntegrand)(out, t)
 end
 
 function (S::AdjointSensitivityIntegrand)(t)
-    # out = similar(S.p.tunable)
     out = similar(S.p)
     out .= false
     S(out, t)
@@ -366,9 +359,7 @@ function _adjoint_sensitivities(sol, sensealg::QuadratureAdjoint, alg; t = nothi
             res, err = quadgk(integrand, sol.prob.tspan[1], sol.prob.tspan[2],
                 atol = abstol, rtol = reltol)
         else
-            tunables, _, _ = canonicalize(Tunable(), integrand.p)
-            # res = zero(integrand.p)'
-            res = zero(tunables)'
+            res = zero(integrand.p)'
 
             # handle discrete dgdp contributions
             if dgdp_discrete !== nothing
@@ -530,4 +521,4 @@ function _update_integrand_and_dgrad(res, sensealg::QuadratureAdjoint, cb, integ
     vecjacobian!(dλ, integrand.y, dλ, integrand.p, t, fakeS; dgrad = dgrad)
     res .-= dgrad
     return integrand
-end
+end

test/parameter_handling.jl

@@ -12,7 +12,7 @@ x = ones(Float32, 2, 3)
 
 nlprob(u, p) = first(model_nls(u, p, st_nls)) .- u
 
-prob = NonlinearProblem(nlprob, zeros(2, 3), ps)
+prob = NonlinearProblem(nlprob, zeros(2, 3), ca)
 
 @test_nowarn solve(prob, NewtonRaphson())
 

test/runtests.jl

@@ -81,13 +81,13 @@ end
 
     if GROUP == "All" || GROUP == "Core6"
         @testset "Core 6" begin
-            # @time @safetestset "Enzyme Closures" include("enzyme_closure.jl")
-            # @time @safetestset "Complex Matrix FiniteDiff Adjoint" include("complex_matrix_finitediff.jl")
-            # @time @safetestset "Null Parameters" include("null_parameters.jl")
-            # @time @safetestset "Forward Mode Prob Kwargs" include("forward_prob_kwargs.jl")
+            @time @safetestset "Enzyme Closures" include("enzyme_closure.jl")
+            @time @safetestset "Complex Matrix FiniteDiff Adjoint" include("complex_matrix_finitediff.jl")
+            @time @safetestset "Null Parameters" include("null_parameters.jl")
+            @time @safetestset "Forward Mode Prob Kwargs" include("forward_prob_kwargs.jl")
             @time @safetestset "Steady State Adjoint" include("steady_state.jl")
-            # @time @safetestset "Concrete Solve Derivatives of Second Order ODEs" include("second_order_odes.jl")
-            # @time @safetestset "Parameter Compatibility Errors" include("parameter_compatibility_errors.jl")
+            @time @safetestset "Concrete Solve Derivatives of Second Order ODEs" include("second_order_odes.jl")
+            @time @safetestset "Parameter Compatibility Errors" include("parameter_compatibility_errors.jl")
         end
     end