fix: halley's method for matrices

Author: avik-pal

lib/SimpleNonlinearSolve/Project.toml

@@ -1,7 +1,7 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "2.2.1"
+version = "2.2.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

lib/SimpleNonlinearSolve/src/halley.jl

@@ -74,7 +74,8 @@ function SciMLBase.__solve(
         end
 
         aᵢ = J_fact \ NLBUtils.safe_vec(fx)
-        hvvp = Utils.compute_hvvp(prob, autodiff, fx_cache, x, aᵢ)
+        hvvp = Utils.compute_hvvp(
+            prob, autodiff, fx_cache, x, NLBUtils.restructure(x, aᵢ))
         bᵢ = J_fact \ NLBUtils.safe_vec(hvvp)
 
         cᵢ_ = NLBUtils.safe_vec(cᵢ)

lib/SimpleNonlinearSolve/src/utils.jl

@@ -171,7 +171,7 @@ function compute_hvvp(prob, autodiff, fx, x, dir)
     else
         @closure (u, p) -> only(DI.pushforward(prob.f, autodiff, u, (dir,), Constant(p)))
     end
-    only(DI.pushforward(jvp_fn, autodiff, x, (dir,), Constant(prob.p)))
+    return only(DI.pushforward(jvp_fn, autodiff, x, (dir,), Constant(prob.p)))
 end
 
 end