fix: minor test fixes

SciML · Sep 27, 2024 · 59347e3 · 59347e3
1 parent be0f8da
commit 59347e3
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 8 deletions.
diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl
@@ -101,7 +101,8 @@ function __internal_init(
         args...; stats, internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN}
     T = promote_type(eltype(fu), eltype(u))
     if u isa Number
-        autodiff = get_concrete_forward_ad(alg.autodiff, prob; check_forward_mode = true)
+        autodiff = get_dense_ad(get_concrete_forward_ad(
+            alg.autodiff, prob; check_forward_mode = true))
         if !(autodiff isa AutoForwardDiff ||
              autodiff isa AutoPolyesterForwardDiff ||
              autodiff isa AutoFiniteDiff)

diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl
@@ -99,10 +99,15 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
                 # While this is technically wasteful, it gives out the type of the Jacobian
                 # which is needed to create the linear solver cache
                 stats.njacs += 1
-                if iip
-                    DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p))
+                if has_analytic_jac
+                    __similar(
+                        fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
                 else
-                    DI.jacobian(f, autodiff, u, Constant(p))
+                    if iip
+                        DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p))
+                    else
+                        DI.jacobian(f, autodiff, u, Constant(p))
+                    end
                 end
             else
                 zero(init_jacobian(sdifft_extras; preserve_immutable = Val(true)))
@@ -120,9 +125,10 @@ function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats,
         autodiff = nothing, kwargs...) where {F}
     fu = f(u, p)
     if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
-        return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing)
+        return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing, nothing)
     end
-    autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
+    autodiff = get_dense_ad(get_concrete_forward_ad(
+        autodiff, prob; check_forward_mode = false))
     di_extras = DI.prepare_derivative(f, get_dense_ad(autodiff), u, Constant(prob.p))
     return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, di_extras, nothing)
 end

diff --git a/test/misc/bruss_tests.jl b/test/misc/bruss_tests.jl
@@ -73,6 +73,6 @@
 
     cache = init(
         prob_brusselator_2d, NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff())))
-    @test maximum(cache.sdifft_extras.jac_cache.coloring.colorvec) == 12
-    @test cache.sdifft_extras.autodiff isa AutoSparse{<:AutoForwardDiff}
+    @test maximum(cache.jac_cache.sdifft_extras.coloring.colorvec) == 12
+    @test cache.jac_cache.autodiff isa AutoSparse{<:AutoForwardDiff}
 end