From 1e7686113573b6085b6be2bedf79ca6d6ad6b23f Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Tue, 24 Oct 2023 11:52:50 -0400
Subject: [PATCH] Generic _axpy!

---
 src/broyden.jl                  | 2 +-
 src/klement.jl                  | 2 +-
 src/lbroyden.jl                 | 2 +-
 src/raphson.jl                  | 2 +-
 src/utils.jl                    | 6 ++++++
 test/nonlinear_least_squares.jl | 4 ++--
 6 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/broyden.jl b/src/broyden.jl
index 9ebd5e4f8..a6981a668 100644
--- a/src/broyden.jl
+++ b/src/broyden.jl
@@ -78,7 +78,7 @@ function perform_step!(cache::GeneralBroydenCache{true})
 
     mul!(_vec(du), J⁻¹, -_vec(fu))
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(α, du, u)
+    _axpy!(α, du, u)
     f(fu2, u, p)
 
     cache.internalnorm(fu2) < cache.abstol && (cache.force_stop = true)
diff --git a/src/klement.jl b/src/klement.jl
index c878faa3b..a16ed2873 100644
--- a/src/klement.jl
+++ b/src/klement.jl
@@ -115,7 +115,7 @@ function perform_step!(cache::GeneralKlementCache{true})
 
     # Line Search
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(α, du, u)
+    _axpy!(α, du, u)
     f(cache.fu2, u, p)
 
     cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true)
diff --git a/src/lbroyden.jl b/src/lbroyden.jl
index bafd04887..efaa1bb04 100644
--- a/src/lbroyden.jl
+++ b/src/lbroyden.jl
@@ -93,7 +93,7 @@ function perform_step!(cache::LimitedMemoryBroydenCache{true})
     T = eltype(u)
 
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(α, du, u)
+    _axpy!(α, du, u)
     f(cache.fu2, u, p)
 
     cache.internalnorm(cache.fu2) < cache.abstol && (cache.force_stop = true)
diff --git a/src/raphson.jl b/src/raphson.jl
index db9e9e322..4a2b53404 100644
--- a/src/raphson.jl
+++ b/src/raphson.jl
@@ -97,7 +97,7 @@ function perform_step!(cache::NewtonRaphsonCache{true})
 
     # Line Search
     α = perform_linesearch!(cache.lscache, u, du)
-    axpy!(-α, du, u)
+    _axpy!(-α, du, u)
     f(cache.fu1, u, p)
 
     cache.internalnorm(fu1) < cache.abstol && (cache.force_stop = true)
diff --git a/src/utils.jl b/src/utils.jl
index 40f04ea4c..afcb68922 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -260,7 +260,13 @@ _try_factorize_and_check_singular!(::Nothing, x) = _issingular(x), false
 _reshape(x, args...) = reshape(x, args...)
 _reshape(x::Number, args...) = x
 
+@generated function _axpy!(α, x, y)
+    hasmethod(axpy!, Tuple{α, x, y}) && return :(axpy!(α, x, y))
+    return :(@. y += α * x)
+end
+
 # Needs Square Matrix
+# FIXME: Remove once https://github.com/SciML/LinearSolve.jl/pull/400 is merged and tagged
 """
     needs_square_A(alg)
 
diff --git a/test/nonlinear_least_squares.jl b/test/nonlinear_least_squares.jl
index c09c9f4d9..c7a02dc58 100644
--- a/test/nonlinear_least_squares.jl
+++ b/test/nonlinear_least_squares.jl
@@ -29,9 +29,9 @@ prob_iip = NonlinearLeastSquaresProblem(NonlinearFunction(loss_function;
 nlls_problems = [prob_oop, prob_iip]
 solvers = [
     GaussNewton(),
-    GaussNewton(; linsolve = CholeskyFactorization()),
+    GaussNewton(; linsolve = LUFactorization()),
     LevenbergMarquardt(),
-    LevenbergMarquardt(; linsolve = CholeskyFactorization()),
+    LevenbergMarquardt(; linsolve = LUFactorization()),
     LeastSquaresOptimJL(:lm),
     LeastSquaresOptimJL(:dogleg),
 ]