Generic _axpy!

SciML · Oct 24, 2023 · 1e76861 · 1e76861
1 parent 30f7873
commit 1e76861
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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),
 ]