add reused_jacobian function

Author: oscardssmith

lib/NonlinearSolveBase/src/NonlinearSolveBase.jl

@@ -75,7 +75,7 @@ include("forward_diff.jl")
 # public for NonlinearSolve.jl and subpackages to use
 @compat(public, (InternalAPI, supports_line_search, supports_trust_region, set_du!))
 @compat(public, (construct_linear_solver, needs_square_A, needs_concrete_A))
-@compat(public, (construct_jacobian_cache,))
+@compat(public, (construct_jacobian_cache, reused_jacobian))
 @compat(public,
     (assert_extension_supported_termination_condition,
     construct_extension_function_wrapper, construct_extension_jac))

lib/NonlinearSolveBase/src/jacobian.jl

@@ -129,6 +129,9 @@ end
 (cache::JacobianCache{<:JacobianOperator})(::Nothing) = error("Please report a bug to NonlinearSolve.jl")
 (cache::JacobianCache)(::Nothing) = error("Please report a bug to NonlinearSolve.jl")
 
+reused_jacobian(cache::JacobianCache, u) = cache.J
+reused_jacobian(cache::JacobianCache{<:JacobianOperator}, u) = StatefulJacobianOperator(cache.J, u, cache.p)
+
 # Core Computation
 ## Numbers
 function (cache::JacobianCache{<:Number})(u)

lib/NonlinearSolveBase/src/wrappers.jl

@@ -106,5 +106,5 @@ function construct_extension_jac(
 
     initial_jacobian isa Val{false} && return J_final
 
-    return J_final, Jₚ(u0)
+    return J_final, J = reused_jacobian(Jₚ, u0)
 end

lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl

@@ -21,7 +21,7 @@ using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm,
                           Utils, InternalAPI, get_timer_output, @static_timeit,
                           update_trace!, L2_NORM, NonlinearSolvePolyAlgorithm,
                           NewtonDescent, DampedNewtonDescent, GeodesicAcceleration,
-                          Dogleg, NonlinearSolveForwardDiffCache
+                          Dogleg, NonlinearSolveForwardDiffCache, reused_jacobian
 using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode,
                  NonlinearFunction,
                  NonlinearLeastSquaresProblem, NonlinearProblem, NoSpecialize

lib/NonlinearSolveFirstOrder/src/solve.jl

@@ -170,7 +170,7 @@ function SciMLBase.__init(
             prob, alg, prob.f, fu, u, prob.p;
             stats, alg.autodiff, linsolve, alg.jvp_autodiff, alg.vjp_autodiff
         )
-        J = jac_cache(u)
+        J = reused_jacobian(cache.jac_cache, cache.u)
 
         descent_cache = InternalAPI.init(
             prob, alg.descent, J, fu, u; stats, abstol, reltol, internalnorm,
@@ -238,7 +238,7 @@ function InternalAPI.step!(
             J = cache.jac_cache(cache.u)
             new_jacobian = true
         else
-            J = cache.jac_cache(cache.u)
+            J = reused_jacobian(cache.jac_cache, cache.u)
             new_jacobian = false
         end
     end

lib/NonlinearSolveQuasiNewton/src/NonlinearSolveQuasiNewton.jl

@@ -18,7 +18,7 @@ using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm,
                           AbstractApproximateJacobianUpdateRule, AbstractDescentDirection,
                           AbstractApproximateJacobianUpdateRuleCache,
                           Utils, InternalAPI, get_timer_output, @static_timeit,
-                          update_trace!, L2_NORM, NewtonDescent
+                          update_trace!, L2_NORM, NewtonDescent, reused_jacobian
 using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode,
                  NonlinearProblem, NonlinearFunction, NoSpecialize
 using SciMLOperators: AbstractSciMLOperator

lib/NonlinearSolveQuasiNewton/src/initialization.jl

@@ -125,7 +125,7 @@ function InternalAPI.init(
     jac_cache = NonlinearSolveBase.construct_jacobian_cache(
         prob, solver, prob.f, fu, u, p; stats, autodiff, linsolve
     )
-    J = alg.structure(jac_cache(u))
+    J = alg.structure(reused_jacobian(jac_cache, u))
     return InitializedApproximateJacobianCache(
         J, alg.structure, alg, jac_cache, false, internalnorm
     )