diff --git a/src/gauss_adjoint.jl b/src/gauss_adjoint.jl index f98075978..aaf71b61f 100644 --- a/src/gauss_adjoint.jl +++ b/src/gauss_adjoint.jl @@ -490,7 +490,7 @@ function (S::GaussIntegrand)(out, t, λ) y = sol(t) end vec_pjac!(out, λ, y, t, S) - out .*= -1 + out = recursive_neg!(out) if S.dgdp !== nothing S.dgdp(dgdp_cache, y, p, t) out .+= dgdp_cache diff --git a/test/adjoint.jl b/test/adjoint.jl index 52fb4bee0..7b101d418 100644 --- a/test/adjoint.jl +++ b/test/adjoint.jl @@ -1063,15 +1063,17 @@ _, res = adjoint_sensitivities(sol_singular_mm, alg, t = ts, reference_sol = ForwardDiff.gradient( p -> G(p, prob_singular_mm, ts, sol -> sum(last, sol.u)), vec(p)) + for salg in [ - QuadratureAdjoint(), - InterpolatingAdjoint(), - BacksolveAdjoint(), - GaussAdjoint() -] + QuadratureAdjoint(), + InterpolatingAdjoint(), + BacksolveAdjoint(), + GaussAdjoint() + ] _, res = adjoint_sensitivities(sol_singular_mm, alg, t = ts, dgdu_discrete = dg_singular, abstol = 1e-14, reltol = 1e-14, sensealg = salg, maxiters = Int(1e6)) + @test res'≈reference_sol rtol=1e-7 end diff --git a/test/parameter_handling.jl b/test/parameter_handling.jl index 0a0f435c5..acd7cd86b 100644 --- a/test/parameter_handling.jl +++ b/test/parameter_handling.jl @@ -1,31 +1,29 @@ using SciMLSensitivity, Lux, Random, Zygote, NonlinearSolve, OrdinaryDiffEq, Test -@static if VERSION ≥ v"1.9" - @info "Testing Nonlinear Solve Adjoint with Nested Parameters" +@info "Testing Nonlinear Solve Adjoint with Nested Parameters" - const model_nls = Chain(Dense(2 => 2, tanh), Dense(2 => 2)) - ps, st = Lux.setup(Random.default_rng(), model_nls) - const st_nls = st +const model_nls = Chain(Dense(2 => 2, tanh), Dense(2 => 2)) +ps, st = Lux.setup(Random.default_rng(), model_nls) +const st_nls = st - x = ones(Float32, 2, 3) +x = ones(Float32, 2, 3) - nlprob(u, p) = first(model_nls(u, p, st_nls)) .- u +nlprob(u, p) = first(model_nls(u, p, st_nls)) .- u - prob = NonlinearProblem(nlprob, zeros(2, 3), ps) +prob = NonlinearProblem(nlprob, zeros(2, 3), ps) - @test_nowarn solve(prob, NewtonRaphson()) +@test_nowarn solve(prob, NewtonRaphson()) - gs = only(Zygote.gradient(ps) do ps - prob = NonlinearProblem(nlprob, zero.(x), ps) - sol = solve(prob, NewtonRaphson()) - return sum(sol.u) - end) +gs = only(Zygote.gradient(ps) do ps + prob = NonlinearProblem(nlprob, zero.(x), ps) + sol = solve(prob, NewtonRaphson()) + return sum(sol.u) +end) - @test gs.layer_1.weight !== nothing - @test gs.layer_1.bias !== nothing - @test gs.layer_2.weight !== nothing - @test gs.layer_2.bias !== nothing -end +@test gs.layer_1.weight !== nothing +@test gs.layer_1.bias !== nothing +@test gs.layer_2.weight !== nothing +@test gs.layer_2.bias !== nothing @info "Testing Gauss Adjoint with Nested Parameters" @@ -43,7 +41,7 @@ prob = ODEProblem(odeprob, ones(2, 3), (0.0f0, 1.0f0), ps) gs = only(Zygote.gradient(ps) do ps prob = ODEProblem(odeprob, ones(2, 3), (0.0f0, 1.0f0), ps) - sol = solve(prob, Tsit5(); sensealg = GaussAdjoint(; autodiff = ZygoteVJP())) + sol = solve(prob, Tsit5(); sensealg = GaussAdjoint(; autojacvec = ZygoteVJP())) return sum(last(sol.u)) end)