Skip to content

Commit

Permalink
test: adjoints
Browse files Browse the repository at this point in the history
  • Loading branch information
avik-pal committed Oct 7, 2024
1 parent fcca2f7 commit 0cbc2fc
Showing 1 changed file with 21 additions and 0 deletions.
21 changes: 21 additions & 0 deletions lib/SimpleNonlinearSolve/test/core/adjoint_tests.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
@testitem "Simple Adjoint Test" tags=[:adjoint] begin
using ForwardDiff, ReverseDiff, SciMLSensitivity, Tracker, Zygote

ff(u, p) = u .^ 2 .- p

function solve_nlprob(p)
prob = NonlinearProblem{false}(ff, [1.0, 2.0], p)
sol = solve(prob, SimpleNewtonRaphson())
res = sol isa AbstractArray ? sol : sol.u
return sum(abs2, res)
end

p = [3.0, 2.0]

∂p_zygote = only(Zygote.gradient(solve_nlprob, p))
∂p_forwarddiff = ForwardDiff.gradient(solve_nlprob, p)
∂p_tracker = Tracker.data(only(Tracker.gradient(solve_nlprob, p)))
∂p_reversediff = ReverseDiff.gradient(solve_nlprob, p)
@test ∂p_zygote ∂p_tracker ∂p_reversediff
@test ∂p_zygote ∂p_forwarddiff ∂p_tracker ∂p_reversediff
end

0 comments on commit 0cbc2fc

Please sign in to comment.