diff --git a/.github/workflows/CI_NonlinearSolve.yml b/.github/workflows/CI_NonlinearSolve.yml index 6f161a07a..a4c14820e 100644 --- a/.github/workflows/CI_NonlinearSolve.yml +++ b/.github/workflows/CI_NonlinearSolve.yml @@ -71,7 +71,7 @@ jobs: GROUP: ${{ matrix.group }} - uses: julia-actions/julia-processcoverage@v1 with: - directories: src,ext + directories: src,ext,lib/SciMLJacobianOperators/src - uses: codecov/codecov-action@v4 with: file: lcov.info diff --git a/.github/workflows/CI_SciMLJacobianOperators.yml b/.github/workflows/CI_SciMLJacobianOperators.yml index 25ba9c652..92daf23e9 100644 --- a/.github/workflows/CI_SciMLJacobianOperators.yml +++ b/.github/workflows/CI_SciMLJacobianOperators.yml @@ -54,7 +54,7 @@ jobs: shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/SciMLJacobianOperators {0} - uses: julia-actions/julia-processcoverage@v1 with: - directories: src,ext + directories: lib/SciMLJacobianOperators/src - uses: codecov/codecov-action@v4 with: file: lcov.info diff --git a/lib/SciMLJacobianOperators/Project.toml b/lib/SciMLJacobianOperators/Project.toml index d5d265d90..5b91928c1 100644 --- a/lib/SciMLJacobianOperators/Project.toml +++ b/lib/SciMLJacobianOperators/Project.toml @@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" @@ -16,23 +17,40 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" [compat] ADTypes = "1.8.1" +Aqua = "0.8.7" ConcreteStructs = "0.2.3" ConstructionBase = "1.5" DifferentiationInterface = "0.5" +Enzyme = "0.12, 0.13" +EnzymeCore = "0.7, 0.8" +ExplicitImports = "1.9.0" FastClosures = "0.3.2" +FiniteDiff = "2.24.0" +ForwardDiff = "0.10.36" InteractiveUtils = "<0.0.1, 1" LinearAlgebra = "1.10" +ReverseDiff = "1.15" SciMLBase = "2.54.0" SciMLOperators = "0.3" Setfield = "1" Test = "1.10" TestItemRunner = "1" +Tracker = "0.2.35" +Zygote = "0.6.70" julia = "1.10" [extras] +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" +FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a" +Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["InteractiveUtils", "Test", "TestItemRunner"] +test = ["Aqua", "Enzyme", "ExplicitImports", "FiniteDiff", "ForwardDiff", "InteractiveUtils", "ReverseDiff", "Test", "TestItemRunner", "Tracker", "Zygote"] diff --git a/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl b/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl index 262443175..3699e71bb 100644 --- a/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl +++ b/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl @@ -1,9 +1,10 @@ module SciMLJacobianOperators -using ADTypes: ADTypes, AutoSparse +using ADTypes: ADTypes, AutoSparse, AutoEnzyme using ConcreteStructs: @concrete using ConstructionBase: ConstructionBase using DifferentiationInterface: DifferentiationInterface +using EnzymeCore: EnzymeCore using FastClosures: @closure using LinearAlgebra: LinearAlgebra using SciMLBase: SciMLBase, AbstractNonlinearProblem, AbstractNonlinearFunction @@ -115,10 +116,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_autodiff = set_function_as_const(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_autodiff = set_function_as_const(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) @@ -259,6 +260,9 @@ end function Base.:*(JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray) return JᵀJ.vjp_operator * (JᵀJ.jvp_operator * x) end +function Base.:*(JᵀJ::StatefulJacobianNormalFormOperator, x::Number) + return JᵀJ.vjp_operator * (JᵀJ.jvp_operator * x) +end function LinearAlgebra.mul!( JᵀJx::AbstractArray, JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray) @@ -284,7 +288,7 @@ function prepare_vjp(::Val{false}, prob::AbstractNonlinearProblem, jac_cache = similar(u, eltype(fu), length(fu), length(u)) return @closure (vJ, v, u, p) -> begin f.jac(jac_cache, u, p) - mul!(vec(vJ), jac_cache', vec(v)) + LinearAlgebra.mul!(vec(vJ), jac_cache', vec(v)) return end return vjp_op, vjp_extras @@ -298,6 +302,8 @@ function prepare_vjp(::Val{false}, prob::AbstractNonlinearProblem, # TODO: Once DI supports const params we can use `p` fₚ = SciMLBase.JacobianWrapper{SciMLBase.isinplace(f)}(f, prob.p) if SciMLBase.isinplace(f) + @assert DI.check_twoarg(autodiff) "Backend: $(autodiff) doesn't support in-place \ + problems." fu_cache = copy(fu) v_fake = copy(fu) di_extras = DI.prepare_pullback(fₚ, fu_cache, autodiff, u, v_fake) @@ -326,11 +332,11 @@ function prepare_jvp(::Val{false}, prob::AbstractNonlinearProblem, jac_cache = similar(u, eltype(fu), length(fu), length(u)) return @closure (Jv, v, u, p) -> begin f.jac(jac_cache, u, p) - mul!(vec(Jv), jac_cache, vec(v)) + LinearAlgebra.mul!(vec(Jv), jac_cache, vec(v)) return end else - return @closure (v, u, p, _) -> reshape(f.jac(u, p) * vec(v), size(u)) + return @closure (v, u, p) -> reshape(f.jac(u, p) * vec(v), size(u)) end end @@ -339,6 +345,8 @@ function prepare_jvp(::Val{false}, prob::AbstractNonlinearProblem, # TODO: Once DI supports const params we can use `p` fₚ = SciMLBase.JacobianWrapper{SciMLBase.isinplace(f)}(f, prob.p) if SciMLBase.isinplace(f) + @assert DI.check_twoarg(autodiff) "Backend: $(autodiff) doesn't support in-place \ + problems." fu_cache = copy(fu) di_extras = DI.prepare_pushforward(fₚ, fu_cache, autodiff, u, u) return @closure (Jv, v, u, p) -> begin @@ -375,6 +383,12 @@ function get_dense_ad(ad::AutoSparse) return dense_ad end +# In our case we know that it is safe to mark the function as const +set_function_as_const(ad) = ad +function set_function_as_const(ad::AutoEnzyme{M, Nothing}) where {M} + return AutoEnzyme(; ad.mode, function_annotation = EnzymeCore.Const) +end + export JacobianOperator, VecJacOperator, JacVecOperator export StatefulJacobianOperator export StatefulJacobianNormalFormOperator diff --git a/lib/SciMLJacobianOperators/test/core_tests.jl b/lib/SciMLJacobianOperators/test/core_tests.jl new file mode 100644 index 000000000..e3b595221 --- /dev/null +++ b/lib/SciMLJacobianOperators/test/core_tests.jl @@ -0,0 +1,274 @@ +@testitem "Scalar Ops" begin + using ADTypes, SciMLBase + using Enzyme, Zygote, ForwardDiff, FiniteDiff, ReverseDiff, Tracker + using SciMLJacobianOperators + + reverse_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Reverse), + AutoZygote(), + AutoReverseDiff(), + AutoTracker(), + AutoFiniteDiff() + ] + + forward_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Forward), + AutoForwardDiff(), + AutoFiniteDiff() + ] + + prob = NonlinearProblem(NonlinearFunction{false}((u, p) -> u^2 - p), 1.0, 2.0) + + analytic_jac(u, p) = 2 * u + analytic_jvp(v, u, p) = 2 * u * v + analytic_vjp(v, u, p) = analytic_jvp(v, u, p) + + @testset "AutoDiff" begin + @testset for jvp_autodiff in forward_ADs, vjp_autodiff in reverse_ADs + jac_op = JacobianOperator(prob, -1.0, 1.0; jvp_autodiff, vjp_autodiff) + + @testset for u in rand(4), v in rand(4) + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + end + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u^2 - p; jvp = analytic_jvp, vjp = analytic_vjp), + 1.0, 2.0) + + @testset "Analytic JVP/VJP" begin + jac_op = JacobianOperator(prob, -1.0, 1.0) + + @testset for u in rand(4), v in rand(4) + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u^2 - p; jac = (u, p) -> 2 * u), 1.0, 2.0) + + @testset "Analytic Jacobian" begin + jac_op = JacobianOperator(prob, -1.0, 1.0) + + @testset for u in rand(4), v in rand(4) + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈2 * u * v atol=1e-5 + @test (sop' * v)≈2 * u * v atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end +end + +@testitem "Inplace Problems" begin + using ADTypes, SciMLBase + using Enzyme, ForwardDiff, FiniteDiff, ReverseDiff + using SciMLJacobianOperators + + reverse_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Reverse), + AutoReverseDiff(), + AutoFiniteDiff() + ] + + forward_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Forward), + AutoForwardDiff(), + AutoFiniteDiff() + ] + + prob = NonlinearProblem( + NonlinearFunction{true}((du, u, p) -> du .= u .^ 2 .- p .+ u[2] * u[1]), [1.0, 3.0], 2.0) + + analytic_jac(u, p) = [2 * u[1]+u[2] u[1]; u[2] 2 * u[2]+u[1]] + analytic_jvp(v, u, p) = analytic_jac(u, p) * v + analytic_vjp(v, u, p) = analytic_jac(u, p)' * v + + analytic_jac!(J, u, p) = (J .= analytic_jac(u, p)) + analytic_jvp!(J, v, u, p) = (J .= analytic_jvp(v, u, p)) + analytic_vjp!(J, v, u, p) = (J .= analytic_vjp(v, u, p)) + + @testset "AutoDiff" begin + @testset for jvp_autodiff in forward_ADs, vjp_autodiff in reverse_ADs + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0; jvp_autodiff, vjp_autodiff) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + end + + prob = NonlinearProblem( + NonlinearFunction{true}((du, u, p) -> du .= u .^ 2 .- p .+ u[2] * u[1]; + jvp = analytic_jvp!, vjp = analytic_vjp!), [1.0, 3.0], 2.0) + + @testset "Analytic JVP/VJP" begin + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + + prob = NonlinearProblem( + NonlinearFunction{true}((du, u, p) -> du .= u .^ 2 .- p .+ u[2] * u[1]; + jac = analytic_jac!), [1.0, 3.0], 2.0) + + @testset "Analytic Jacobian" begin + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end +end + +@testitem "Out-of-place Problems" begin + using ADTypes, SciMLBase + using Enzyme, ForwardDiff, FiniteDiff, ReverseDiff, Zygote, Tracker + using SciMLJacobianOperators + + reverse_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Reverse), + AutoZygote(), + AutoTracker(), + AutoReverseDiff(), + AutoFiniteDiff() + ] + + forward_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Forward), + AutoForwardDiff(), + AutoFiniteDiff() + ] + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u .^ 2 .- p .+ u[2] * u[1]), [1.0, 3.0], 2.0) + + analytic_jac(u, p) = [2 * u[1]+u[2] u[1]; u[2] 2 * u[2]+u[1]] + analytic_jvp(v, u, p) = analytic_jac(u, p) * v + analytic_vjp(v, u, p) = analytic_jac(u, p)' * v + + @testset "AutoDiff" begin + @testset for jvp_autodiff in forward_ADs, vjp_autodiff in reverse_ADs + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0; jvp_autodiff, vjp_autodiff) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + end + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u .^ 2 .- p .+ u[2] * u[1]; + vjp = analytic_vjp, jvp = analytic_jvp), [1.0, 3.0], 2.0) + + @testset "Analytic JVP/VJP" begin + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u .^ 2 .- p .+ u[2] * u[1]; + jac = analytic_jac), [1.0, 3.0], 2.0) + + @testset "Analytic Jacobian" begin + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end +end + +@testitem "Aqua" begin + using SciMLJacobianOperators, Aqua + + Aqua.test_all(SciMLJacobianOperators) +end + +@testitem "Explicit Imports" begin + using SciMLJacobianOperators, ExplicitImports + + @test check_no_implicit_imports(SciMLJacobianOperators) === nothing + @test check_no_stale_explicit_imports(SciMLJacobianOperators) === nothing + @test check_all_qualified_accesses_via_owners(SciMLJacobianOperators) === nothing +end