diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index 822eb72..df39e47 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -25,9 +25,33 @@ jobs: version: - '1' - 'lts' - - 'pre' + # 'pre' skipped due to ReTestItems incompatibility with Julia 1.13 + # See: https://github.com/SciML/LineSearch.jl/issues/47 uses: "SciML/.github/.github/workflows/tests.yml@v1" with: julia-version: "${{ matrix.version }}" group: "all" secrets: "inherit" + + enzyme: + name: "Enzyme Tests" + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - 'lts' + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/cache@v2 + - name: "Run Enzyme tests" + run: | + julia --project=test/enzyme -e ' + using Pkg + Pkg.develop(PackageSpec(path = pwd())) + Pkg.instantiate() + ' + julia --project=test/enzyme test/enzyme/runtests.jl diff --git a/Project.toml b/Project.toml index 85d689a..b2eb00e 100644 --- a/Project.toml +++ b/Project.toml @@ -26,7 +26,6 @@ ADTypes = "1.9" CommonSolve = "0.2.4" ConcreteStructs = "0.2.3" DifferentiationInterface = "0.6.2, 0.7" -Enzyme = "0.13.3" ExplicitImports = "1.14.0" FastClosures = "0.3" FiniteDiff = "2.24.0" @@ -50,7 +49,6 @@ julia = "1.10" [extras] DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" @@ -65,4 +63,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["DifferentiationInterface", "Enzyme", "ExplicitImports", "FiniteDiff", "ForwardDiff", "Hwloc", "InteractiveUtils", "LineSearches", "NonlinearProblemLibrary", "ReTestItems", "ReverseDiff", "Test", "Tracker", "Zygote"] +test = ["DifferentiationInterface", "ExplicitImports", "FiniteDiff", "ForwardDiff", "Hwloc", "InteractiveUtils", "LineSearches", "NonlinearProblemLibrary", "ReTestItems", "ReverseDiff", "Test", "Tracker", "Zygote"] diff --git a/test/custom_optimizer_tests.jl b/test/custom_optimizer_tests.jl index df8fbbe..b922dcb 100644 --- a/test/custom_optimizer_tests.jl +++ b/test/custom_optimizer_tests.jl @@ -1,4 +1,5 @@ # Test based on https://julianlsolvers.github.io/LineSearches.jl/stable/examples/generated/customoptimizer.html +# Note: Enzyme tests are in a separate test group (test/enzyme/) @testsetup module CustomOptimizer using LinearAlgebra, SciMLBase, LineSearch, SciMLJacobianOperators @@ -51,7 +52,7 @@ end @testitem "LineSearches.jl: Custom Optimizer" setup = [CustomOptimizer] begin using LineSearches, SciMLBase - using ADTypes, Tracker, ForwardDiff, Zygote, Enzyme, ReverseDiff, FiniteDiff + using ADTypes, Tracker, ForwardDiff, Zygote, ReverseDiff, FiniteDiff @testset "OOP Problem" begin nlf(x, p) = [p[1] - x[1], 10.0 * (x[2] - x[1]^2)] @@ -59,7 +60,7 @@ end @testset for autodiff in ( AutoTracker(), AutoForwardDiff(), AutoZygote(), - AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), + AutoReverseDiff(), AutoFiniteDiff(), ) @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), @@ -82,7 +83,7 @@ end nlp = NonlinearProblem(nlf, [-1.0, 1.0], [1.0]) @testset for autodiff in ( - AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), + AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff(), ) @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), @@ -103,7 +104,7 @@ end @testitem "Native Line Search: Custom Optimizer" setup = [CustomOptimizer] begin using SciMLBase - using ADTypes, Tracker, ForwardDiff, Zygote, Enzyme, ReverseDiff, FiniteDiff + using ADTypes, Tracker, ForwardDiff, Zygote, ReverseDiff, FiniteDiff @testset "OOP Problem" begin nlf(x, p) = [p[1] - x[1], 10.0 * (x[2] - x[1]^2)] @@ -111,7 +112,7 @@ end @testset for autodiff in ( AutoTracker(), AutoForwardDiff(), AutoZygote(), - AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), + AutoReverseDiff(), AutoFiniteDiff(), ) @testset "method: $(nameof(typeof(method)))" for method in ( LiFukushimaLineSearch(), @@ -133,7 +134,7 @@ end nlp = NonlinearProblem(nlf, [-1.0, 1.0], [1.0]) @testset for autodiff in ( - AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), + AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff(), ) @testset "method: $(nameof(typeof(method)))" for method in ( LiFukushimaLineSearch(), diff --git a/test/enzyme/Project.toml b/test/enzyme/Project.toml new file mode 100644 index 0000000..fb64b5b --- /dev/null +++ b/test/enzyme/Project.toml @@ -0,0 +1,30 @@ +[deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" +ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[compat] +ADTypes = "1.9" +DifferentiationInterface = "0.6.2, 0.7" +Enzyme = "0.13.3" +FiniteDiff = "2.24.0" +ForwardDiff = "0.10.36, 1" +LineSearch = "0.1" +LineSearches = "7.3.0" +ReverseDiff = "1.15.3" +SciMLBase = "2.53.1" +SciMLJacobianOperators = "0.1" +Test = "1.10" +Tracker = "0.2.35" +Zygote = "0.6.71" diff --git a/test/enzyme/runtests.jl b/test/enzyme/runtests.jl new file mode 100644 index 0000000..9920a7d --- /dev/null +++ b/test/enzyme/runtests.jl @@ -0,0 +1,297 @@ +using Test, LineSearch, SciMLBase, ADTypes, Enzyme +using DifferentiationInterface, ForwardDiff, Tracker, Zygote, ReverseDiff, FiniteDiff +using SciMLJacobianOperators, LinearAlgebra, LineSearches + +@info "Running Enzyme tests with Julia $(VERSION)" + +# Test setup module for custom optimizer tests +module CustomOptimizer + using LinearAlgebra, SciMLBase, LineSearch, SciMLJacobianOperators + + function gradient_descent( + prob, alg; g_atol::Real = 1.0e-5, maxiters::Int = 10000, autodiff = nothing + ) + u = copy(prob.u0) + if SciMLBase.isinplace(prob) + fu = similar(u) + prob.f(fu, u, prob.p) + else + fu = prob.f(u, prob.p) + end + + ls_cache = init(prob, alg, fu, u) + vjp_op = VecJacOperator(prob, fu, u; autodiff) + + alphas = Float64[] + iter = 0 + for _ in 1:maxiters + iter += 1 + svjp_op = StatefulJacobianOperator(vjp_op, u, prob.p) + gs = svjp_op * fu .* 2 + δu = -gs + + ls_sol = solve!(ls_cache, u, δu) + + push!(alphas, ls_sol.step_size) + @. u = u + ls_sol.step_size * δu + gnorm = norm(gs) + + if SciMLBase.isinplace(prob) + fu = similar(u) + prob.f(fu, u, prob.p) + else + fu = prob.f(u, prob.p) + end + + LineSearch.callback_into_cache!(ls_cache, fu) + + gnorm < g_atol && break + end + + return fu, u, iter, alphas + end + + export gradient_descent +end + +using .CustomOptimizer + +# Test setup module for root finding tests +module RootFinding + using SciMLBase, DifferentiationInterface, ForwardDiff + using SciMLBase: AbstractNonlinearProblem + const DI = DifferentiationInterface + + function newton_raphson(prob::AbstractNonlinearProblem, ls) + if SciMLBase.isinplace(prob) + return newton_raphson_iip(prob, ls) + else + return newton_raphson_oop(prob, ls) + end + end + + function newton_raphson_oop(prob::AbstractNonlinearProblem, ls) + u = copy(prob.u0) + fu = prob.f(u, prob.p) + + ls_cache = init(prob, ls, fu, u) + + alphas = Float64[] + iter = 0 + for _ in 1:100 + iter += 1 + + maximum(abs, fu) < 1.0e-8 && return true, fu, u, iter, alphas + + J = DI.jacobian(prob.f, AutoForwardDiff(), u, Constant(prob.p)) + δu = -J \ fu + + ls_sol = solve!(ls_cache, u, δu) + + push!(alphas, ls_sol.step_size) + @. u = u + ls_sol.step_size * δu + + fu = prob.f(u, prob.p) + end + + return false, fu, u, iter, alphas + end + + function newton_raphson_iip(prob::AbstractNonlinearProblem, ls) + u = copy(prob.u0) + fu = similar(u) + fu2 = similar(u) + prob.f(fu, u, prob.p) + + ls_cache = init(prob, ls, fu, u) + + alphas = Float64[] + iter = 0 + for _ in 1:100 + iter += 1 + + maximum(abs, fu) < 1.0e-8 && return true, fu, u, iter, alphas + + J = DI.jacobian(prob.f, fu2, AutoForwardDiff(), u, Constant(prob.p)) + δu = -J \ fu + + ls_sol = solve!(ls_cache, u, δu) + + push!(alphas, ls_sol.step_size) + @. u = u + ls_sol.step_size * δu + + prob.f(fu, u, prob.p) + end + + return false, fu, u, iter, alphas + end + + export newton_raphson +end + +using .RootFinding + +const OOP_AUTODIFFS = (AutoEnzyme(),) +const IIP_AUTODIFFS = (AutoEnzyme(),) + +@testset "Enzyme: LineSearches.jl Custom Optimizer" begin + @testset "OOP Problem" begin + nlf(x, p) = [p[1] - x[1], 10.0 * (x[2] - x[1]^2)] + nlp = NonlinearProblem(nlf, [-1.0, 1.0], [1.0]) + + @testset for autodiff in OOP_AUTODIFFS + @testset "method: $(nameof(typeof(method)))" for method in ( + LineSearches.BackTracking(; order = 3), + StrongWolfe(), + HagerZhang(), + MoreThuente(), + ) + linesearch = LineSearchesJL(; method, autodiff) + fu, u, iter, alphas = gradient_descent(nlp, linesearch; autodiff) + + @test fu ≈ [0.0, 0.0] atol = 1.0e-2 + @test u ≈ [1.0, 1.0] atol = 1.0e-2 + @test !all(isone, alphas) + end + end + end + + @testset "In-Place Problem" begin + nlf(dx, x, p) = (dx .= [p[1] - x[1], 10.0 * (x[2] - x[1]^2)]) + nlp = NonlinearProblem(nlf, [-1.0, 1.0], [1.0]) + + @testset for autodiff in IIP_AUTODIFFS + @testset "method: $(nameof(typeof(method)))" for method in ( + LineSearches.BackTracking(; order = 3), + StrongWolfe(), + HagerZhang(), + MoreThuente(), + ) + linesearch = LineSearchesJL(; method, autodiff) + fu, u, iter, alphas = gradient_descent(nlp, linesearch; autodiff) + + @test fu ≈ [0.0, 0.0] atol = 1.0e-2 + @test u ≈ [1.0, 1.0] atol = 1.0e-2 + @test !all(isone, alphas) + end + end + end +end + +@testset "Enzyme: Native Line Search Custom Optimizer" begin + @testset "OOP Problem" begin + nlf(x, p) = [p[1] - x[1], 10.0 * (x[2] - x[1]^2)] + nlp = NonlinearProblem(nlf, [-1.0, 1.0], [1.0]) + + @testset for autodiff in OOP_AUTODIFFS + @testset "method: $(nameof(typeof(method)))" for method in ( + LineSearch.BackTracking(; order = Val(3), autodiff), + LineSearch.BackTracking(; order = Val(2), autodiff), + ) + fu, u, iter, alphas = gradient_descent(nlp, method; autodiff) + + @test fu ≈ [0.0, 0.0] atol = 1.0e-1 + @test u ≈ [1.0, 1.0] atol = 1.0e-1 + @test !all(isone, alphas) + end + end + end + + @testset "In-Place Problem" begin + nlf(dx, x, p) = (dx .= [p[1] - x[1], 10.0 * (x[2] - x[1]^2)]) + nlp = NonlinearProblem(nlf, [-1.0, 1.0], [1.0]) + + @testset for autodiff in IIP_AUTODIFFS + @testset "method: $(nameof(typeof(method)))" for method in ( + LineSearch.BackTracking(; order = Val(3), autodiff), + LineSearch.BackTracking(; order = Val(2), autodiff), + ) + fu, u, iter, alphas = gradient_descent(nlp, method; autodiff) + + @test fu ≈ [0.0, 0.0] atol = 1.0e-1 + @test u ≈ [1.0, 1.0] atol = 1.0e-1 + @test !all(isone, alphas) + end + end + end +end + +@testset "Enzyme: LineSearches.jl Newton Raphson" begin + @testset "OOP Problem" begin + nlf(x, p) = x .^ 2 .- p + nlp = NonlinearProblem(nlf, [-1.0, 1.0], [3.0]) + + @testset for autodiff in OOP_AUTODIFFS + @testset "method: $(nameof(typeof(method)))" for method in ( + LineSearches.BackTracking(; order = 3), + StrongWolfe(), + HagerZhang(), + MoreThuente(), + Static(), + ) + linesearch = LineSearchesJL(; method, autodiff) + converged, fu, u, iter, alphas = newton_raphson(nlp, linesearch) + + @test fu ≈ [0.0, 0.0] atol = 1.0e-3 + @test abs.(u) ≈ sqrt.([3.0, 3.0]) atol = 1.0e-3 + end + end + end + + @testset "In-Place Problem" begin + nlf(dx, x, p) = (dx .= x .^ 2 .- p) + nlp = NonlinearProblem(nlf, [-1.0, 1.0], [3.0]) + + @testset for autodiff in IIP_AUTODIFFS + @testset "method: $(nameof(typeof(method)))" for method in ( + LineSearches.BackTracking(; order = 3), + StrongWolfe(), + HagerZhang(), + MoreThuente(), + Static(), + ) + linesearch = LineSearchesJL(; method, autodiff) + converged, fu, u, iter, alphas = newton_raphson(nlp, linesearch) + + @test fu ≈ [0.0, 0.0] atol = 1.0e-3 + @test abs.(u) ≈ sqrt.([3.0, 3.0]) atol = 1.0e-3 + end + end + end +end + +@testset "Enzyme: Native Line Search Newton Raphson" begin + @testset "OOP Problem" begin + nlf(x, p) = x .^ 2 .- p + nlp = NonlinearProblem(nlf, [-1.0, 1.0], [3.0]) + + @testset for autodiff in OOP_AUTODIFFS + @testset "method: $(nameof(typeof(method)))" for method in ( + LineSearch.BackTracking(; order = Val(3), autodiff), + LineSearch.BackTracking(; order = Val(2), autodiff), + ) + converged, fu, u, iter, alphas = newton_raphson(nlp, method) + + @test fu ≈ [0.0, 0.0] atol = 1.0e-3 + @test abs.(u) ≈ sqrt.([3.0, 3.0]) atol = 1.0e-3 + end + end + end + + @testset "In-Place Problem" begin + nlf(dx, x, p) = (dx .= x .^ 2 .- p) + nlp = NonlinearProblem(nlf, [-1.0, 1.0], [3.0]) + + @testset for autodiff in IIP_AUTODIFFS + @testset "method: $(nameof(typeof(method)))" for method in ( + LineSearch.BackTracking(; order = Val(3), autodiff), + LineSearch.BackTracking(; order = Val(2), autodiff), + ) + converged, fu, u, iter, alphas = newton_raphson(nlp, method) + + @test fu ≈ [0.0, 0.0] atol = 1.0e-3 + @test abs.(u) ≈ sqrt.([3.0, 3.0]) atol = 1.0e-3 + end + end + end +end diff --git a/test/root_finding_tests.jl b/test/root_finding_tests.jl index 4231e85..8b543ad 100644 --- a/test/root_finding_tests.jl +++ b/test/root_finding_tests.jl @@ -1,5 +1,6 @@ # Here we write out Newton Raphson and test integration with LineSearch.jl. Main tests are # over at NonlinearSolve.jl and SimpleNonlinearSolve.jl +# Note: Enzyme tests are in a separate test group (test/enzyme/) @testsetup module RootFinding using SciMLBase, DifferentiationInterface, ForwardDiff @@ -76,7 +77,7 @@ end @testitem "LineSearches.jl: Newton Raphson" setup = [RootFinding] begin using LineSearches, SciMLBase - using ADTypes, Tracker, ForwardDiff, Zygote, Enzyme, ReverseDiff, FiniteDiff + using ADTypes, Tracker, ForwardDiff, Zygote, ReverseDiff, FiniteDiff @testset "OOP Problem" begin nlf(x, p) = x .^ 2 .- p @@ -84,7 +85,7 @@ end @testset for autodiff in ( AutoTracker(), AutoForwardDiff(), AutoZygote(), - AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), + AutoReverseDiff(), AutoFiniteDiff(), ) @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), @@ -107,7 +108,7 @@ end nlp = NonlinearProblem(nlf, [-1.0, 1.0], [3.0]) @testset for autodiff in ( - AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), + AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff(), ) @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), @@ -128,7 +129,7 @@ end @testitem "Native Line Search: Newton Raphson" setup = [RootFinding] begin using SciMLBase - using ADTypes, Tracker, ForwardDiff, Zygote, Enzyme, ReverseDiff, FiniteDiff + using ADTypes, Tracker, ForwardDiff, Zygote, ReverseDiff, FiniteDiff @testset "OOP Problem" begin nlf(x, p) = x .^ 2 .- p @@ -146,7 +147,7 @@ end @testset for autodiff in ( AutoTracker(), AutoForwardDiff(), AutoZygote(), - AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), + AutoReverseDiff(), AutoFiniteDiff(), ) @testset "method: $(nameof(typeof(method)))" for method in ( BackTracking(; order = Val(3), autodiff), @@ -175,7 +176,7 @@ end end @testset for autodiff in ( - AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), + AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff(), ) @testset "method: $(nameof(typeof(method)))" for method in ( BackTracking(; order = Val(3), autodiff),