Skip to content

Commit

Permalink
test: fix core testing
Browse files Browse the repository at this point in the history
  • Loading branch information
avik-pal committed Sep 24, 2024
1 parent 9466b78 commit 37c56f7
Show file tree
Hide file tree
Showing 6 changed files with 34 additions and 10 deletions.
9 changes: 8 additions & 1 deletion lib/SciMLJacobianOperators/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"]
21 changes: 18 additions & 3 deletions lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
module SciMLJacobianOperators

using ADTypes: ADTypes, AutoSparse
using ConcreteStructs: @concrete
using ConstructionBase: ConstructionBase
using DifferentiationInterface: DifferentiationInterface
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions lib/SciMLJacobianOperators/test/runtests.jl
Original file line number Diff line number Diff line change
@@ -1 +1,5 @@
using TestItemRunner, InteractiveUtils

@info sprint(InteractiveUtils.versioninfo)

@run_package_tests
2 changes: 1 addition & 1 deletion src/internal/jacobian.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
6 changes: 3 additions & 3 deletions test/core/rootfind_tests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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]),
Expand Down
2 changes: 0 additions & 2 deletions test/misc/qa_tests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 37c56f7

Please sign in to comment.