Handle Klement

Author: avik-pal

src/nlsolve/klement.jl

@@ -25,7 +25,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleKlement, args...;
     @bb δx² = similar(x)
 
     for _ in 1:maxiters
-        any(iszero, J) && (J = __init_identity_jacobian!!(J))
+        any(any(iszero, J)) && (J = __init_identity_jacobian!!(J))
 
         @bb @. δx = fprev / J
 

src/utils.jl

@@ -22,6 +22,9 @@ __max_tdir(a, b, x0, x1) = ifelse(x1 > x0, max(a, b), min(a, b))
 __standard_tag(::Nothing, f::F, x::AbstractArray{T}) where {F, T} = ForwardDiff.Tag(f, T)
 __standard_tag(tag::ForwardDiff.Tag, f::F, x::AbstractArray{T}) where {F, T} = tag
 __standard_tag(tag, f::F, x::AbstractArray{T}) where {F, T} = ForwardDiff.Tag(tag, T)
+__standard_tag(::Nothing, f::F, x::T) where {F, T <: Number} = ForwardDiff.Tag(f, T)
+__standard_tag(tag::ForwardDiff.Tag, f::F, x::T) where {F, T <: Number} = tag
+__standard_tag(tag, f::F, x::T) where {F, T <: Number} = ForwardDiff.Tag(tag, T)
 
 __pick_forwarddiff_chunk(x) = ForwardDiff.Chunk(length(x))
 function __pick_forwarddiff_chunk(x::StaticArray)
@@ -40,7 +43,7 @@ function __get_jacobian_config(ad::AutoForwardDiff{CS}, f::F, x) where {F, CS}
 end
 function __get_jacobian_config(ad::AutoForwardDiff{CS}, f!::F, y, x) where {F, CS}
     ck = (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) : ForwardDiff.Chunk{CS}()
-    tag = __standard_tag(ad.tag, f, x)
+    tag = __standard_tag(ad.tag, f!, x)
     return ForwardDiff.JacobianConfig(f!, y, x, ck, tag)
 end