From 02cf007362c3af97042f2fd90a6b9b5db226c6f7 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    | 5 +++++
 5 files changed, 9 insertions(+), 4 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..7ea62cc4e 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -260,6 +260,11 @@ _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{typeof(α), typeof(x), typeof(y)}) && return :(axpy!(α, x, y))
+    return :(y .= α * x)
+end
+
 # Needs Square Matrix
 """
     needs_square_A(alg)