Fix forward AD

Author: avik-pal

src/ad.jl

@@ -1,40 +1,61 @@
 function scalar_nlsolve_ad(prob, alg, args...; kwargs...)
     f = prob.f
     p = value(prob.p)
-
     u0 = value(prob.u0)
     newprob = NonlinearProblem(f, u0, p; prob.kwargs...)
 
     sol = solve(newprob, alg, args...; kwargs...)
 
     uu = sol.u
-    if p isa Number
-        f_p = ForwardDiff.derivative(Base.Fix1(f, uu), p)
-    else
-        f_p = ForwardDiff.gradient(Base.Fix1(f, uu), p)
-    end
+    f_p = scalar_nlsolve_∂f_∂p(f, uu, p)
+    f_x = scalar_nlsolve_∂f_∂u(f, uu, p)
+
+    z_arr = -inv(f_x) * f_p
 
-    f_x = ForwardDiff.derivative(Base.Fix2(f, p), uu)
     pp = prob.p
-    sumfun = let f_x′ = -f_x
-        ((fp, p),) -> (fp / f_x′) * ForwardDiff.partials(p)
+    sumfun = ((z, p),) -> [zᵢ * ForwardDiff.partials(p) for zᵢ in z]
+    if uu isa Number
+        partials = sum(sumfun, zip(z_arr, pp))
+    else
+        partials = sum(sumfun, zip(eachcol(z_arr), pp))
     end
-    partials = sum(sumfun, zip(f_p, pp))
+
     return sol, partials
 end
 
-function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, SVector}, iip,
-        <:Dual{T, V, P}}, alg::AbstractNewtonAlgorithm, args...;
+function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, SVector, <:AbstractArray},
+        iip, <:Dual{T, V, P}}, alg::AbstractNewtonAlgorithm, args...;
     kwargs...) where {iip, T, V, P}
     sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-    return SciMLBase.build_solution(prob, alg, Dual{T, V, P}(sol.u, partials), sol.resid;
-        sol.retcode)
+    dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
+    return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode)
 end
 
-function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, SVector}, iip,
-        <:AbstractArray{<:Dual{T, V, P}}}, alg::AbstractNewtonAlgorithm, args...;
+function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, SVector, <:AbstractArray},
+        iip, <:AbstractArray{<:Dual{T, V, P}}}, alg::AbstractNewtonAlgorithm, args...;
     kwargs...) where {iip, T, V, P}
     sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-    return SciMLBase.build_solution(prob, alg, Dual{T, V, P}(sol.u, partials), sol.resid;
-        sol.retcode)
+    dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
+    return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode)
+end
+
+function scalar_nlsolve_∂f_∂p(f, u, p)
+    ff = p isa Number ? ForwardDiff.derivative :
+         (u isa Number ? ForwardDiff.gradient : ForwardDiff.jacobian)
+    return ff(Base.Fix1(f, u), p)
+end
+
+function scalar_nlsolve_∂f_∂u(f, u, p)
+    ff = u isa Number ? ForwardDiff.derivative : ForwardDiff.jacobian
+    return ff(Base.Fix2(f, p), u)
+end
+
+function scalar_nlsolve_dual_soln(u::Number, partials,
+    ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}}) where {T, V, P}
+    return Dual{T, V, P}(u, partials[1])
+end
+
+function scalar_nlsolve_dual_soln(u::AbstractArray, partials,
+    ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}}) where {T, V, P}
+    return map(((uᵢ, pᵢ),) -> Dual{T, V, P}(uᵢ, pᵢ), zip(u, partials))
 end

test/basictests.jl

@@ -57,15 +57,13 @@ end
         @test (@ballocated solve!($cache)) ≤ 64
     end
 
-    # FIXME: Even the previous tests were broken, but due to a typo in the tests they
-    #        accidentally passed
     @testset "[OOP] [Immutable AD] p: $(p)" for p in 1.0:0.1:100.0
         @test begin
             res = benchmark_nlsolve_oop(quadratic_f, @SVector[1.0, 1.0], p)
             res_true = sqrt(p)
             all(res.u .≈ res_true)
         end
-        @test_broken ForwardDiff.derivative(p -> benchmark_nlsolve_oop(quadratic_f,
+        @test ForwardDiff.derivative(p -> benchmark_nlsolve_oop(quadratic_f,
             @SVector[1.0, 1.0], p).u[end], p) ≈ 1 / (2 * sqrt(p))
     end
 
@@ -101,11 +99,9 @@ end
     @test nlprob_iterator_interface(quadratic_f, p, Val(false)) ≈ sqrt.(p)
     @test nlprob_iterator_interface(quadratic_f!, p, Val(true)) ≈ sqrt.(p)
 
-    probN = NonlinearProblem(quadratic_f, @SVector[1.0, 1.0], 2.0)
-    @testset "ADType: $(autodiff) u0: $(u0)" for autodiff in (false, true,
+    @testset "ADType: $(autodiff) u0: $(_nameof(u0))" for autodiff in (false, true,
             AutoSparseForwardDiff(), AutoSparseFiniteDiff(), AutoZygote(),
-            AutoSparseZygote(),
-            AutoSparseEnzyme()), u0 in (1.0, [1.0, 1.0], @SVector[1.0, 1.0])
+            AutoSparseZygote(), AutoSparseEnzyme()), u0 in (1.0, [1.0, 1.0])
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
         @test all(solve(probN, NewtonRaphson(; autodiff)).u .≈ sqrt(2.0))
     end
@@ -149,8 +145,6 @@ end
         @test (@ballocated solve!($cache)) ≤ 64
     end
 
-    # FIXME: Even the previous tests were broken, but due to a typo in the tests they
-    #        accidentally passed
     @testset "[OOP] [Immutable AD] radius_update_scheme: $(radius_update_scheme) p: $(p)" for radius_update_scheme in radius_update_schemes,
         p in 1.0:0.1:100.0
 
@@ -160,7 +154,7 @@ end
             res_true = sqrt(p)
             all(res.u .≈ res_true)
         end
-        @test_broken ForwardDiff.derivative(p -> benchmark_nlsolve_oop(quadratic_f,
+        @test ForwardDiff.derivative(p -> benchmark_nlsolve_oop(quadratic_f,
             @SVector[1.0, 1.0], p; radius_update_scheme).u[end], p) ≈ 1 / (2 * sqrt(p))
     end
 
@@ -204,11 +198,9 @@ end
     @test nlprob_iterator_interface(quadratic_f, p, Val(false)) ≈ sqrt.(p)
     @test nlprob_iterator_interface(quadratic_f!, p, Val(true)) ≈ sqrt.(p)
 
-    probN = NonlinearProblem(quadratic_f, @SVector[1.0, 1.0], 2.0)
-    @testset "ADType: $(autodiff) u0: $(u0) radius_update_scheme: $(radius_update_scheme)" for autodiff in (false,
+    @testset "ADType: $(autodiff) u0: $(_nameof(u0)) radius_update_scheme: $(radius_update_scheme)" for autodiff in (false,
             true, AutoSparseForwardDiff(), AutoSparseFiniteDiff(), AutoZygote(),
-            AutoSparseZygote(), AutoSparseEnzyme()),
-        u0 in (1.0, [1.0, 1.0], @SVector[1.0, 1.0]),
+            AutoSparseZygote(), AutoSparseEnzyme()), u0 in (1.0, [1.0, 1.0]),
         radius_update_scheme in radius_update_schemes
 
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
@@ -302,15 +294,13 @@ end
         @test (@ballocated solve!($cache)) ≤ 64
     end
 
-    # FIXME: Even the previous tests were broken, but due to a typo in the tests they
-    #        accidentally passed
     @testset "[OOP] [Immutable AD] p: $(p)" for p in 1.0:0.1:100.0
         @test begin
             res = benchmark_nlsolve_oop(quadratic_f, @SVector[1.0, 1.0], p)
             res_true = sqrt(p)
             all(res.u .≈ res_true)
         end
-        @test_broken ForwardDiff.derivative(p -> benchmark_nlsolve_oop(quadratic_f,
+        @test ForwardDiff.derivative(p -> benchmark_nlsolve_oop(quadratic_f,
             @SVector[1.0, 1.0], p).u[end], p) ≈ 1 / (2 * sqrt(p))
     end
 
@@ -330,11 +320,9 @@ end
     @test ForwardDiff.jacobian(p -> [benchmark_nlsolve_oop(quadratic_f2, 0.5, p).u], p) ≈
           ForwardDiff.jacobian(t, p)
 
-    probN = NonlinearProblem(quadratic_f, @SVector[1.0, 1.0], 2.0)
-    @testset "ADType: $(autodiff) u0: $(u0)" for autodiff in (false, true,
+    @testset "ADType: $(autodiff) u0: $(_nameof(u0))" for autodiff in (false, true,
             AutoSparseForwardDiff(), AutoSparseFiniteDiff(), AutoZygote(),
-            AutoSparseZygote(),
-            AutoSparseEnzyme()), u0 in (1.0, [1.0, 1.0], @SVector[1.0, 1.0])
+            AutoSparseZygote(), AutoSparseEnzyme()), u0 in (1.0, [1.0, 1.0])
         probN = NonlinearProblem(quadratic_f, u0, 2.0)
         @test all(solve(probN, LevenbergMarquardt(; autodiff)).u .≈ sqrt(2.0))
     end