fix: don't ignore sparsity for dense_ad

Author: avik-pal

src/NonlinearSolve.jl

@@ -7,8 +7,8 @@ end
 using Reexport: @reexport
 using PrecompileTools: @compile_workload, @setup_workload
 
-using ADTypes: ADTypes, AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff,
-               AutoZygote, AutoEnzyme, AutoSparse
+using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff,
+               AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme, AutoSparse
 # FIXME: deprecated, remove in future
 using ADTypes: AutoSparseFiniteDiff, AutoSparseForwardDiff, AutoSparsePolyesterForwardDiff,
                AutoSparseZygote

src/internal/jacobian.jl

@@ -60,7 +60,7 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
 
     if !has_analytic_jac && needs_jac
         autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
-        sd = __sparsity_detection_alg(f, autodiff)
+        sd = sparsity_detection_alg(f, autodiff)
         sparse_jac = !(sd isa NoSparsityDetection)
         # Eventually we want to do everything via DI. But for now, we just do the dense via DI
         if sparse_jac
@@ -174,9 +174,13 @@ function (cache::JacobianCache{iip})(
     end
 end
 
-# Sparsity Detection Choices
-@inline __sparsity_detection_alg(_, _) = NoSparsityDetection()
-@inline function __sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse)
+function sparsity_detection_alg(f::NonlinearFunction, ad::AbstractADType)
+    # TODO: Also handle case where colorvec is provided
+    f.sparsity === nothing && return NoSparsityDetection()
+    return sparsity_detection_alg(f, AutoSparse(ad; sparsity_detector = f.sparsity))
+end
+
+function sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse)
     if f.sparsity === nothing
         if f.jac_prototype === nothing
             is_extension_loaded(Val(:Symbolics)) && return SymbolicsSparsityDetection()
@@ -200,8 +204,7 @@ end
     end
 
     if SciMLBase.has_colorvec(f)
-        return PrecomputedJacobianColorvec(; jac_prototype,
-            f.colorvec,
+        return PrecomputedJacobianColorvec(; jac_prototype, f.colorvec,
             partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode)
     else
         return JacPrototypeSparsityDetection(; jac_prototype)