diff --git a/src/klement.jl b/src/klement.jl
index 0db04dc6e..ed6a3b172 100644
--- a/src/klement.jl
+++ b/src/klement.jl
@@ -174,9 +174,9 @@ function perform_step!(cache::GeneralKlementCache{iip, IJ}) where {iip, IJ}
             cache.J_cache = jacobian!!(cache.J_cache, cache)
             cache.J = __get_diagonal!!(cache.J, cache.J_cache)
         end
-        ill_conditioned = __is_ill_conditioned(cache.J)
+        ill_conditioned = __is_ill_conditioned(_vec(cache.J))
     elseif IJ === :identity
-        ill_conditioned = __is_ill_conditioned(cache.J)
+        ill_conditioned = __is_ill_conditioned(_vec(cache.J))
     end
 
     if ill_conditioned
@@ -228,7 +228,7 @@ function perform_step!(cache::GeneralKlementCache{iip, IJ}) where {iip, IJ}
         @bb @. cache.J += ((cache.fu - cache.fu_cache_2 - cache.J * cache.du) /
                            ifelse(iszero(cache.Jdu), T(1e-5), cache.Jdu)) * cache.du *
                           (cache.J^2)
-    else
+    elseif IJ === :true_jacobian
         # Klement Updates to the Full Jacobian don't work for most problems, we should
         # probably be using the Broyden Update Rule here
         @bb @. cache.J_cache = cache.J'^2
@@ -241,6 +241,8 @@ function perform_step!(cache::GeneralKlementCache{iip, IJ}) where {iip, IJ}
         @bb @. cache.J_cache *= cache.J
         @bb cache.J_cache_2 = cache.J_cache × cache.J
         @bb cache.J .+= cache.J_cache_2
+    else
+        error("Invalid `init_jacobian` value")
     end
 
     @bb copyto!(cache.fu_cache_2, cache.fu)
diff --git a/src/utils.jl b/src/utils.jl
index 44c255bd4..a781a954d 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -551,6 +551,9 @@ end
     end
     return J
 end
+@inline function __get_diagonal!!(J::AbstractArray, J_full::AbstractMatrix)
+    return _restructure(J, __get_diagonal!!(_vec(J), J_full))
+end
 @inline __get_diagonal!!(J::Number, J_full::Number) = J_full
 
 @inline __diag(x::AbstractMatrix) = diag(x)
diff --git a/test/matrix_resizing.jl b/test/matrix_resizing.jl
index 59a537ace..c9579d9cf 100644
--- a/test/matrix_resizing.jl
+++ b/test/matrix_resizing.jl
@@ -1,7 +1,7 @@
-using NonlinearSolve, Test
+using NonlinearSolve, Test, StableRNGs
 
 ff(u, p) = u .* u .- p
-u0 = rand(2, 2)
+u0 = rand(StableRNG(0), 2, 2)
 p = 2.0
 vecprob = NonlinearProblem(ff, vec(u0), p)
 prob = NonlinearProblem(ff, u0, p)
@@ -13,7 +13,7 @@ for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransien
 end
 
 fiip(du, u, p) = (du .= u .* u .- p)
-u0 = rand(2, 2)
+u0 = rand(StableRNG(0), 2, 2)
 p = 2.0
 vecprob = NonlinearProblem(fiip, vec(u0), p)
 prob = NonlinearProblem(fiip, u0, p)