fix jacobians

Author: oscardssmith

lib/NonlinearSolveBase/src/jacobian.jl

@@ -61,7 +61,7 @@ function construct_jacobian_cache(
     end
 
     J = if !needs_jac
-        StatefulJacobianOperator(JacobianOperator(prob, fu, u; jvp_autodiff, vjp_autodiff), u, p)
+        JacobianOperator(prob, fu, u; jvp_autodiff, vjp_autodiff)
     else
         if f.jac_prototype === nothing
             # While this is technically wasteful, it gives out the type of the Jacobian
@@ -96,7 +96,7 @@ function construct_jacobian_cache(
         linsolve = missing
 )
     if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
-        return JacobianCache(u, f, fu, u, p, stats, autodiff, nothing)
+        return JacobianCache(fu, f, fu, p, stats, autodiff, nothing)
     end
     if autodiff === nothing
         throw(ArgumentError("`autodiff` argument to `construct_jacobian_cache` must be \
@@ -124,10 +124,12 @@ function InternalAPI.reinit!(cache::JacobianCache; p = cache.p, kwargs...)
     cache.p = p
 end
 
-# Core Computation
-(cache::JacobianCache)(::Nothing) = cache.J
-(cache::JacobianCache{<:Number})(::Nothing) = cache.J
+# Deprecations
+(cache::JacobianCache{<:Number})(::Nothing) = error("Please report a bug to NonlinearSolve.jl")
+(cache::JacobianCache{<:JacobianOperator})(::Nothing) = error("Please report a bug to NonlinearSolve.jl")
+(cache::JacobianCache)(::Nothing) = error("Please report a bug to NonlinearSolve.jl")
 
+# Core Computation
 ## Numbers
 function (cache::JacobianCache{<:Number})(u)
     cache.stats.njacs += 1
@@ -168,6 +170,10 @@ function (cache::JacobianCache)(u)
     end
 end
 
+function (cache::JacobianCache{<:JacobianOperator})(u)
+    return StatefulJacobianOperator(cache.J, u, cache.p)
+end
+
 # Sparse Automatic Differentiation
 function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType)
     if f.sparsity === nothing

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ₚ(nothing)
+    return J_final, Jₚ(u0)
 end

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(nothing)
+        J = jac_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(nothing)
+            J = cache.jac_cache(cache.u)
             new_jacobian = false
         end
     end

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(nothing))
+    J = alg.structure(jac_cache(u))
     return InitializedApproximateJacobianCache(
         J, alg.structure, alg, jac_cache, false, internalnorm
     )