Patch tracing and LM

Author: avik-pal

src/levenberg.jl

@@ -273,7 +273,7 @@ function perform_step!(cache::LevenbergMarquardtCache{iip, fastls}) where {iip,
             b = cache.rhs_tmp, linu = _vec(cache.v), cache.p, reltol = cache.abstol)
         @bb @. cache.v = -linres.u
     else
-        @bb cache.u_cache_2 = transpose(J) × cache.fu
+        @bb cache.u_cache_2 = transpose(cache.J) × cache.fu
         @bb @. cache.mat_tmp = cache.JᵀJ + cache.λ * cache.DᵀD
         linres = dolinsolve(alg.precs, linsolve; A = cache.mat_tmp,
             b = _vec(cache.u_cache_2), linu = _vec(cache.v), cache.p, reltol = cache.abstol)
@@ -288,7 +288,7 @@ function perform_step!(cache::LevenbergMarquardtCache{iip, fastls}) where {iip,
     @bb @. cache.u_cache_2 = cache.u + cache.h * cache.v
     evaluate_f(cache, cache.u_cache_2, cache.p, Val(:fu_cache_2))
 
-    # The following lines do: cache.a = -J \ cache.fu_tmp
+    # The following lines do: cache.a = -cache.mat_tmp \ cache.fu_tmp
     # NOTE: Don't pass `A` in again, since we want to reuse the previous solve
     @bb cache.Jv = cache.J × cache.v
     @bb @. cache.fu_cache_2 = (2 / cache.h) *
@@ -332,7 +332,7 @@ function perform_step!(cache::LevenbergMarquardtCache{iip, fastls}) where {iip,
                 @bb @. cache.fu = cache.fu_cache_2 - cache.fu
                 check_and_update!(cache.tc_cache_2, cache, cache.fu, cache.u, cache.u_cache)
             end
-            @bb copyto!(cache.fu_cache, cache.fu_cache_2)
+            @bb copyto!(cache.fu, cache.fu_cache_2)
             @bb copyto!(cache.v_cache, cache.v)
             cache.norm_v_old = norm_v
             cache.loss_old = loss

src/trace.jl

@@ -209,8 +209,8 @@ function update_trace!(trace::NonlinearSolveTrace{ShT, StT}, iter, u, fu, J, δu
         return trace
     end
 
-    show_now = ShT && (iter % trace.trace_level.print_frequency == 1)
-    store_now = StT && (iter % trace.trace_level.store_frequency == 1)
+    show_now = ShT && (mod1(iter, trace.trace_level.print_frequency) == 1)
+    store_now = StT && (mod1(iter, trace.trace_level.store_frequency) == 1)
     (show_now || store_now) && (entry = __trace_entry(trace.trace_level, iter, u, fu, J,
         δu, α))
     store_now && push!(trace.history, entry)
@@ -230,8 +230,8 @@ function update_trace_with_invJ!(trace::NonlinearSolveTrace{ShT, StT}, iter, u,
         return trace
     end
 
-    show_now = ShT && (iter % trace.trace_level.print_frequency == 1)
-    store_now = StT && (iter % trace.trace_level.store_frequency == 1)
+    show_now = ShT && (mod1(iter, trace.trace_level.print_frequency) == 1)
+    store_now = StT && (mod1(iter, trace.trace_level.store_frequency) == 1)
     if show_now || store_now
         J_ = trace.trace_level isa TraceMinimal ? J : inv(J)
         entry = __trace_entry(trace.trace_level, iter, u, fu, J_, δu, α)

src/utils.jl

@@ -448,15 +448,16 @@ function __sum_JᵀJ!!(y, J)
     end
 end
 
-function __update_LM_diagonal!!(y::Diagonal, x::AbstractVector)
+@inline __update_LM_diagonal!!(y::Number, x::Number) = max(y, x)
+@inline function __update_LM_diagonal!!(y::Diagonal, x::AbstractVector)
     if setindex_trait(y.diag) === CanSetindex()
         @. y.diag = max(y.diag, x)
         return y
     else
         return Diagonal(max.(y.diag, x))
     end
 end
-@views function __update_LM_diagonal!!(y::Diagonal, x::AbstractMatrix)
+@inline @views function __update_LM_diagonal!!(y::Diagonal, x::AbstractMatrix)
     x_diag = x[diagind(x)]
     if setindex_trait(y.diag) === CanSetindex()
         @. y.diag = max(y.diag, x_diag)