test: fix core testing

Author: avik-pal

lib/SciMLJacobianOperators/Project.toml

@@ -4,6 +4,7 @@ authors = ["Avik Pal <[email protected]> and contributors"]
 version = "0.1.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
@@ -14,18 +15,24 @@ SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 
 [compat]
+ADTypes = "1.8.1"
 ConcreteStructs = "0.2.3"
 ConstructionBase = "1.5"
 DifferentiationInterface = "0.5"
 FastClosures = "0.3.2"
+InteractiveUtils = "<0.0.1, 1"
 LinearAlgebra = "1.10"
+SciMLBase = "2.54.0"
 SciMLOperators = "0.3"
 Setfield = "1"
+Test = "1.10"
+TestItemRunner = "1"
 julia = "1.10"
 
 [extras]
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
 
 [targets]
-test = ["Test", "TestItemRunner"]
+test = ["InteractiveUtils", "Test", "TestItemRunner"]

lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl

@@ -1,5 +1,6 @@
 module SciMLJacobianOperators
 
+using ADTypes: ADTypes, AutoSparse
 using ConcreteStructs: @concrete
 using ConstructionBase: ConstructionBase
 using DifferentiationInterface: DifferentiationInterface
@@ -101,6 +102,7 @@ for op in (:adjoint, :transpose)
         (; output_cache, input_cache) = operator
         @set! operator.output_cache = input_cache
         @set! operator.input_cache = output_cache
+        @set! operator.size = reverse(operator.size)
         return operator
     end
 end
@@ -113,7 +115,10 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff =
     iip = SciMLBase.isinplace(prob)
     T = promote_type(eltype(u), eltype(fu))
 
+    vjp_autodiff = get_dense_ad(vjp_autodiff)
     vjp_op = prepare_vjp(skip_vjp, prob, f, u, fu; autodiff = vjp_autodiff)
+
+    jvp_autodiff = get_dense_ad(jvp_autodiff)
     jvp_op = prepare_jvp(skip_jvp, prob, f, u, fu; autodiff = jvp_autodiff)
 
     output_cache = fu isa Number ? T(fu) : similar(fu, T)
@@ -212,6 +217,7 @@ for op in (:adjoint, :transpose)
 end
 
 Base.:*(J::StatefulJacobianOperator, v::AbstractArray) = J.jac_op(v, J.u, J.p)
+Base.:*(J::StatefulJacobianOperator, v::Number) = J.jac_op(v, J.u, J.p)
 
 function LinearAlgebra.mul!(
         Jv::AbstractArray, J::StatefulJacobianOperator, v::AbstractArray)
@@ -256,8 +262,8 @@ end
 
 function LinearAlgebra.mul!(
         JᵀJx::AbstractArray, JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray)
-    mul!(JᵀJ.cache, JᵀJ.jvp_operator, x)
-    mul!(JᵀJx, JᵀJ.vjp_operator, JᵀJ.cache)
+    LinearAlgebra.mul!(JᵀJ.cache, JᵀJ.jvp_operator, x)
+    LinearAlgebra.mul!(JᵀJx, JᵀJ.vjp_operator, JᵀJ.cache)
     return JᵀJx
 end
 
@@ -313,7 +319,7 @@ end
 
 function prepare_jvp(::Val{false}, prob::AbstractNonlinearProblem,
         f::AbstractNonlinearFunction, u, fu; autodiff = nothing)
-    SciMLBase.has_vjp(f) && return f.vjp
+    SciMLBase.has_jvp(f) && return f.jvp
 
     if autodiff === nothing && SciMLBase.has_jac(f)
         if SciMLBase.isinplace(f)
@@ -360,6 +366,15 @@ function prepare_scalar_op(::Val{false}, prob::AbstractNonlinearProblem,
     return @closure (v, u, p) -> DI.derivative(fₚ, autodiff, u, di_extras) * v
 end
 
+get_dense_ad(::Nothing) = nothing
+get_dense_ad(ad) = ad
+function get_dense_ad(ad::AutoSparse)
+    dense_ad = ADTypes.dense_ad(ad)
+    @warn "Sparse AD backend: $(ad) is being used for VJP/JVP computation. Using the dense \
+           backend: $(dense_ad) instead."
+    return dense_ad
+end
+
 export JacobianOperator, VecJacOperator, JacVecOperator
 export StatefulJacobianOperator
 export StatefulJacobianNormalFormOperator

lib/SciMLJacobianOperators/test/runtests.jl

@@ -1 +1,5 @@
+using TestItemRunner, InteractiveUtils
 
+@info sprint(InteractiveUtils.versioninfo)
+
+@run_package_tests

src/internal/jacobian.jl

@@ -103,10 +103,10 @@ function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats,
     if !(autodiff isa AutoForwardDiff ||
          autodiff isa AutoPolyesterForwardDiff ||
          autodiff isa AutoFiniteDiff)
-        autodiff = AutoFiniteDiff()
         # Other cases are not properly supported so we fallback to finite differencing
         @warn "Scalar AD is supported only for AutoForwardDiff and AutoFiniteDiff. \
                Detected $(autodiff). Falling back to AutoFiniteDiff."
+        autodiff = AutoFiniteDiff()
     end
     return JacobianCache{false}(
         u, f, uf, u, u, p, nothing, alg, stats, autodiff, nothing, nothing)

test/core/rootfind_tests.jl

@@ -125,7 +125,7 @@ end
     u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0)
     linear_solvers = [nothing, LUFactorization(), KrylovJL_GMRES(), \]
 
-    @testset "[OOP] u0: $(typeof(u0)) radius_update_scheme: $(radius_update_scheme) linear_solver: $(linsolve)" for u0 in u0s,
+    @testset "[OOP] u0: $(typeof(u0)) $(radius_update_scheme) $(_nameof(linsolve))" for u0 in u0s,
         radius_update_scheme in radius_update_schemes,
         linsolve in linear_solvers
 
@@ -141,7 +141,7 @@ end
         @test (@ballocated solve!($cache)) < 200
     end
 
-    @testset "[IIP] u0: $(typeof(u0)) radius_update_scheme: $(radius_update_scheme) linear_solver: $(linsolve)" for u0 in ([
+    @testset "[IIP] u0: $(typeof(u0)) $(radius_update_scheme) $(_nameof(linsolve))" for u0 in ([
             1.0, 1.0],),
         radius_update_scheme in radius_update_schemes,
         linsolve in linear_solvers
@@ -162,7 +162,7 @@ end
     @test nlprob_iterator_interface(quadratic_f, p, Val(false), TrustRegion()) ≈ sqrt.(p)
     @test nlprob_iterator_interface(quadratic_f!, p, Val(true), TrustRegion()) ≈ sqrt.(p)
 
-    @testset "ADType: $(autodiff) u0: $(_nameof(u0)) radius_update_scheme: $(radius_update_scheme)" for autodiff in (
+    @testset "$(_nameof(autodiff)) u0: $(_nameof(u0)) $(radius_update_scheme)" for autodiff in (
             AutoSparse(AutoForwardDiff()), AutoSparse(AutoFiniteDiff()),
             AutoZygote(), AutoSparse(AutoZygote()), AutoSparse(AutoEnzyme())),
         u0 in (1.0, [1.0, 1.0]),

test/misc/qa_tests.jl

@@ -24,8 +24,6 @@ end
 
     @test check_no_implicit_imports(NonlinearSolve;
         skip = (NonlinearSolve, Base, Core, SimpleNonlinearSolve, SciMLBase)) === nothing
-
     @test check_no_stale_explicit_imports(NonlinearSolve) === nothing
-
     @test check_all_qualified_accesses_via_owners(NonlinearSolve) === nothing
 end