diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index 822eb72..00f0dfd 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -31,3 +31,26 @@ jobs: 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..f974a75 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" @@ -42,6 +41,7 @@ ReTestItems = "1.28.0" ReverseDiff = "1.15.3" SciMLBase = "2.53.1" SciMLJacobianOperators = "0.1" +StaticArrays = "1.9" StaticArraysCore = "1.4" Test = "1.10" Tracker = "0.2.35" @@ -50,7 +50,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" @@ -60,9 +59,10 @@ LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141" ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" 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", "StaticArrays", "Test", "Tracker", "Zygote"] diff --git a/test/alloc_tests.jl b/test/alloc_tests.jl new file mode 100644 index 0000000..6e43041 --- /dev/null +++ b/test/alloc_tests.jl @@ -0,0 +1,212 @@ +# Allocation tests to prevent performance regressions +# These tests verify that critical IIP (in-place) code paths remain zero-allocation +# Note: Allocation behavior varies across Julia versions, so tests are skipped on older versions + +@testitem "Allocation Tests: LiFukushimaLineSearch IIP" tags = [:alloc] begin + using LineSearch, SciMLBase, CommonSolve, Test + + # Skip allocation tests on Julia < 1.13 where allocation behavior differs + if VERSION < v"1.13" + @test_skip "Skipped on Julia $(VERSION) - allocation tests require Julia 1.13+" + else + function nlf_iip!(dx, x, p) + dx[1] = x[1]^2 - p[1] + dx[2] = x[2]^2 - p[2] + return nothing + end + + nlp = NonlinearProblem(nlf_iip!, [-1.0, 1.0], [3.0, 3.0]) + fu = similar(nlp.u0) + nlf_iip!(fu, nlp.u0, nlp.p) + u = copy(nlp.u0) + du = [-0.5, 0.5] + + alg = LiFukushimaLineSearch() + cache = init(nlp, alg, fu, u) + + # Warmup + for _ in 1:10 + solve!(cache, u, du) + end + + @testset "solve! is zero-allocation" begin + allocs = @allocated solve!(cache, u, du) + @test allocs == 0 + end + end +end + +@testitem "Allocation Tests: RobustNonMonotoneLineSearch IIP" tags = [:alloc] begin + using LineSearch, SciMLBase, CommonSolve, Test + + # Skip allocation tests on Julia < 1.13 where allocation behavior differs + if VERSION < v"1.13" + @test_skip "Skipped on Julia $(VERSION) - allocation tests require Julia 1.13+" + else + function nlf_iip!(dx, x, p) + dx[1] = x[1]^2 - p[1] + dx[2] = x[2]^2 - p[2] + return nothing + end + + nlp = NonlinearProblem(nlf_iip!, [-1.0, 1.0], [3.0, 3.0]) + fu = similar(nlp.u0) + nlf_iip!(fu, nlp.u0, nlp.p) + u = copy(nlp.u0) + du = [-0.5, 0.5] + + alg = RobustNonMonotoneLineSearch() + cache = init(nlp, alg, fu, u) + + # Warmup + for _ in 1:10 + solve!(cache, u, du) + end + + @testset "solve! is zero-allocation" begin + allocs = @allocated solve!(cache, u, du) + @test allocs == 0 + end + + @testset "callback_into_cache! is zero-allocation" begin + allocs = @allocated LineSearch.callback_into_cache!(cache, fu) + @test allocs == 0 + end + end +end + +@testitem "Allocation Tests: NoLineSearch" tags = [:alloc] begin + using LineSearch, SciMLBase, CommonSolve, Test + + # Skip allocation tests on Julia < 1.13 where allocation behavior differs + if VERSION < v"1.13" + @test_skip "Skipped on Julia $(VERSION) - allocation tests require Julia 1.13+" + else + function nlf_iip!(dx, x, p) + dx[1] = x[1]^2 - p[1] + dx[2] = x[2]^2 - p[2] + return nothing + end + + nlp = NonlinearProblem(nlf_iip!, [-1.0, 1.0], [3.0, 3.0]) + fu = similar(nlp.u0) + nlf_iip!(fu, nlp.u0, nlp.p) + u = copy(nlp.u0) + du = [-0.5, 0.5] + + alg = NoLineSearch(1.0) + cache = init(nlp, alg, fu, u) + + # Warmup + for _ in 1:10 + solve!(cache, u, du) + end + + @testset "solve! is zero-allocation" begin + allocs = @allocated solve!(cache, u, du) + @test allocs == 0 + end + end +end + +@testitem "Allocation Tests: StaticLiFukushimaLineSearch" tags = [:alloc] begin + using LineSearch, SciMLBase, CommonSolve, StaticArrays, Test + + # Skip allocation tests on Julia < 1.13 where allocation behavior differs + if VERSION < v"1.13" + @test_skip "Skipped on Julia $(VERSION) - allocation tests require Julia 1.13+" + else + static_f(u::SVector{2}, p) = SVector(u[1]^2 - p[1], u[2]^2 - p[2]) + + prob = NonlinearProblem(static_f, SVector(-1.0, 1.0), SVector(3.0, 3.0)) + fu = static_f(prob.u0, prob.p) + u = prob.u0 + du = SVector(-0.5, 0.5) + + # Use nan_maxiters=nothing to get the static (non-allocating) path + alg = LiFukushimaLineSearch(; nan_maxiters = nothing) + cache = init(prob, alg, fu, u) + + # Warmup + for _ in 1:10 + solve!(cache, u, du) + end + + @testset "solve! is zero-allocation" begin + allocs = @allocated solve!(cache, u, du) + @test allocs == 0 + end + end +end + +@testitem "Allocation Tests: Scalar LiFukushimaLineSearch" tags = [:alloc] begin + using LineSearch, SciMLBase, CommonSolve, Test + + # Skip allocation tests on Julia < 1.13 where allocation behavior differs + if VERSION < v"1.13" + @test_skip "Skipped on Julia $(VERSION) - allocation tests require Julia 1.13+" + else + scalar_f(u, p) = u^2 - p + + prob = NonlinearProblem(scalar_f, 2.0, 4.0) + fu = scalar_f(prob.u0, prob.p) + u = prob.u0 + du = -0.5 + + # Use nan_maxiters=nothing to get the static (non-allocating) path + alg = LiFukushimaLineSearch(; nan_maxiters = nothing) + cache = init(prob, alg, fu, u) + + # Warmup + for _ in 1:10 + solve!(cache, u, du) + end + + @testset "solve! is zero-allocation" begin + allocs = @allocated solve!(cache, u, du) + @test allocs == 0 + end + end +end + +@testitem "Allocation Tests: Larger Problems (10D)" tags = [:alloc] begin + using LineSearch, SciMLBase, CommonSolve, Test + + # Skip allocation tests on Julia < 1.13 where allocation behavior differs + if VERSION < v"1.13" + @test_skip "Skipped on Julia $(VERSION) - allocation tests require Julia 1.13+" + else + function nlf_10d!(dx, x, p) + for i in 1:10 + dx[i] = x[i]^2 - p[i] + end + return nothing + end + + nlp = NonlinearProblem(nlf_10d!, ones(10), 3 * ones(10)) + fu = similar(nlp.u0) + nlf_10d!(fu, nlp.u0, nlp.p) + u = copy(nlp.u0) + du = -0.5 * ones(10) + + @testset "LiFukushimaLineSearch 10D" begin + alg = LiFukushimaLineSearch() + cache = init(nlp, alg, fu, u) + for _ in 1:10 + solve!(cache, u, du) + end + allocs = @allocated solve!(cache, u, du) + @test allocs == 0 + end + + @testset "RobustNonMonotoneLineSearch 10D" begin + alg = RobustNonMonotoneLineSearch() + cache = init(nlp, alg, fu, u) + for _ in 1:10 + solve!(cache, u, du) + end + allocs = @allocated solve!(cache, u, du) + @test allocs == 0 + end + end +end diff --git a/test/custom_optimizer_tests.jl b/test/custom_optimizer_tests.jl index df8fbbe..26c10e1 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,16 +52,19 @@ 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 + + const OOP_AUTODIFFS = ( + AutoTracker(), AutoForwardDiff(), AutoZygote(), + AutoReverseDiff(), AutoFiniteDiff(), + ) + const IIP_AUTODIFFS = (AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff()) @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 ( - AutoTracker(), AutoForwardDiff(), AutoZygote(), - AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), - ) + @testset for autodiff in OOP_AUTODIFFS @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), StrongWolfe(), @@ -81,9 +85,7 @@ end 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 ( - AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), - ) + @testset for autodiff in IIP_AUTODIFFS @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), StrongWolfe(), @@ -103,16 +105,19 @@ 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 + + const OOP_AUTODIFFS = ( + AutoTracker(), AutoForwardDiff(), AutoZygote(), + AutoReverseDiff(), AutoFiniteDiff(), + ) + const IIP_AUTODIFFS = (AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff()) @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 ( - AutoTracker(), AutoForwardDiff(), AutoZygote(), - AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), - ) + @testset for autodiff in OOP_AUTODIFFS @testset "method: $(nameof(typeof(method)))" for method in ( LiFukushimaLineSearch(), NoLineSearch(0.001), @@ -132,9 +137,7 @@ end 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 ( - AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), - ) + @testset for autodiff in IIP_AUTODIFFS @testset "method: $(nameof(typeof(method)))" for method in ( LiFukushimaLineSearch(), NoLineSearch(0.001), 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..6382ed7 --- /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 ( + BackTracking(; order = Val(3), autodiff), + 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 ( + BackTracking(; order = Val(3), autodiff), + 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 ( + BackTracking(; order = Val(3), autodiff), + 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 ( + BackTracking(; order = Val(3), autodiff), + 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..720e93d 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,16 +77,19 @@ 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 + + const OOP_AUTODIFFS = ( + AutoTracker(), AutoForwardDiff(), AutoZygote(), + AutoReverseDiff(), AutoFiniteDiff(), + ) + const IIP_AUTODIFFS = (AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff()) @testset "OOP Problem" begin nlf(x, p) = x .^ 2 .- p nlp = NonlinearProblem(nlf, [-1.0, 1.0], [3.0]) - @testset for autodiff in ( - AutoTracker(), AutoForwardDiff(), AutoZygote(), - AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), - ) + @testset for autodiff in OOP_AUTODIFFS @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), StrongWolfe(), @@ -106,9 +110,7 @@ end nlf(dx, x, p) = (dx .= x .^ 2 .- p) nlp = NonlinearProblem(nlf, [-1.0, 1.0], [3.0]) - @testset for autodiff in ( - AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), - ) + @testset for autodiff in IIP_AUTODIFFS @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), StrongWolfe(), @@ -128,7 +130,13 @@ 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 + + const OOP_AUTODIFFS = ( + AutoTracker(), AutoForwardDiff(), AutoZygote(), + AutoReverseDiff(), AutoFiniteDiff(), + ) + const IIP_AUTODIFFS = (AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff()) @testset "OOP Problem" begin nlf(x, p) = x .^ 2 .- p @@ -144,10 +152,7 @@ end @test abs.(u) ≈ sqrt.([3.0, 3.0]) atol = 1.0e-1 end - @testset for autodiff in ( - AutoTracker(), AutoForwardDiff(), AutoZygote(), - AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), - ) + @testset for autodiff in OOP_AUTODIFFS @testset "method: $(nameof(typeof(method)))" for method in ( BackTracking(; order = Val(3), autodiff), BackTracking(; order = Val(2), autodiff), @@ -174,9 +179,7 @@ end @test abs.(u) ≈ sqrt.([3.0, 3.0]) atol = 1.0e-1 end - @testset for autodiff in ( - AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), - ) + @testset for autodiff in IIP_AUTODIFFS @testset "method: $(nameof(typeof(method)))" for method in ( BackTracking(; order = Val(3), autodiff), BackTracking(; order = Val(2), autodiff),