Forward Mode overloads for Least Squares Problem

Project.toml

@@ -1,7 +1,7 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "1.4.3"
+version = "1.5.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -22,11 +22,13 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 SimpleNonlinearSolveChainRulesCoreExt = "ChainRulesCore"
 SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff"
 SimpleNonlinearSolveStaticArraysExt = "StaticArrays"
+SimpleNonlinearSolveZygoteExt = "Zygote"
 
 [compat]
 ADTypes = "0.2.6"

ext/SimpleNonlinearSolveZygoteExt.jl

@@ -0,0 +1,7 @@
+module SimpleNonlinearSolveZygoteExt
+
+import SimpleNonlinearSolve
+
+SimpleNonlinearSolve.__is_extension_loaded(::Val{:Zygote}) = true
+
+end

src/ad.jl

@@ -24,7 +24,8 @@ for algType in (Bisection, Brent, Alefeld, Falsi, ITP, Ridder)
     end
 end
 
-function __nlsolve_ad(prob, alg, args...; kwargs...)
+function __nlsolve_ad(
+        prob::Union{IntervalNonlinearProblem, NonlinearProblem}, alg, args...; kwargs...)
     p = value(prob.p)
     if prob isa IntervalNonlinearProblem
         tspan = value.(prob.tspan)

test/core/least_squares_tests.jl

@@ -12,6 +12,12 @@
         return ŷ .- y_target
     end
 
+    function loss_function!(resid, θ, p)
+        ŷ = true_function(p, θ)
+        @. resid = ŷ - y_target
+        return
+    end
+
     θ_init = θ_true .+ 0.1
     prob_oop = NonlinearLeastSquaresProblem{false}(loss_function, θ_init, x)
 
@@ -21,4 +27,14 @@
         sol = solve(prob_oop, solver)
         @test norm(sol.resid, Inf) < 1e-12
     end
+
+    prob_iip = NonlinearLeastSquaresProblem(
+        NonlinearFunction{true}(loss_function!, resid_prototype = zeros(length(y_target))), θ_init, x)
+
+    @testset "Solver: $(nameof(typeof(solver)))" for solver in [
+        SimpleNewtonRaphson(AutoForwardDiff()), SimpleGaussNewton(AutoForwardDiff()),
+        SimpleNewtonRaphson(AutoFiniteDiff()), SimpleGaussNewton(AutoFiniteDiff())]
+        sol = solve(prob_iip, solver)
+        @test norm(sol.resid, Inf) < 1e-12
+    end
 end