diff --git a/Project.toml b/Project.toml index decb388..2c089a2 100644 --- a/Project.toml +++ b/Project.toml @@ -32,6 +32,7 @@ FiniteDiff = "2.24.0" ForwardDiff = "0.10.36, 1" Hwloc = "3" InteractiveUtils = "<0.0.1, 1" +JET = "0.9, 0.10, 0.11" LineSearches = "7.3.0" LinearAlgebra = "1.10" MaybeInplace = "0.1.4" @@ -54,6 +55,7 @@ FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141" ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" @@ -63,4 +65,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", "Enzyme", "ExplicitImports", "FiniteDiff", "ForwardDiff", "Hwloc", "InteractiveUtils", "JET", "LineSearches", "NonlinearProblemLibrary", "ReTestItems", "ReverseDiff", "Test", "Tracker", "Zygote"] diff --git a/test/jet_tests.jl b/test/jet_tests.jl new file mode 100644 index 0000000..76c0134 --- /dev/null +++ b/test/jet_tests.jl @@ -0,0 +1,39 @@ +@testitem "JET Static Analysis" tags=[:jet] begin + using JET, LineSearch, SciMLBase, ADTypes, ForwardDiff + using CommonSolve: init, solve! + + f_oop(u, p) = u .^ 2 .- p + prob = NonlinearProblem(f_oop, [1.0, 2.0], [3.0]) + u = [1.0, 2.0] + fu = f_oop(u, prob.p) + du = -fu + + @testset "NoLineSearch" begin + alg = NoLineSearch() + cache = init(prob, alg, fu, u) + @test_opt target_modules = (LineSearch,) solve!(cache, u, du) + end + + @testset "LiFukushimaLineSearch" begin + alg = LiFukushimaLineSearch() + cache = init(prob, alg, fu, u) + @test_opt target_modules = (LineSearch,) solve!(cache, u, du) + end + + @testset "RobustNonMonotoneLineSearch" begin + alg = RobustNonMonotoneLineSearch() + cache = init(prob, alg, fu, u) + @test_opt target_modules = (LineSearch,) solve!(cache, u, du) + end + + @testset "BackTracking" begin + alg = BackTracking(; autodiff = AutoForwardDiff()) + cache = init(prob, alg, fu, u) + # BackTracking uses closures with captured variables from external packages + # (SciMLJacobianOperators), which JET cannot fully analyze for type stability. + # We test for errors only, not optimization/type stability issues. + report = JET.report_call(solve!, (typeof(cache), typeof(u), typeof(du)); + target_modules = (LineSearch,)) + @test length(JET.get_reports(report)) == 0 + end +end diff --git a/test/root_finding_tests.jl b/test/root_finding_tests.jl index 4231e85..f1409b5 100644 --- a/test/root_finding_tests.jl +++ b/test/root_finding_tests.jl @@ -78,14 +78,25 @@ end using LineSearches, SciMLBase using ADTypes, Tracker, ForwardDiff, Zygote, Enzyme, ReverseDiff, FiniteDiff + # Enzyme is exempted from Julia v1 (stable) testing + const ENZYME_COMPAT = VERSION >= v"1.11" + @testset "OOP Problem" begin nlf(x, p) = x .^ 2 .- p nlp = NonlinearProblem(nlf, [-1.0, 1.0], [3.0]) - @testset for autodiff in ( + autodiffs_oop = if ENZYME_COMPAT + ( AutoTracker(), AutoForwardDiff(), AutoZygote(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), ) + else + ( + AutoTracker(), AutoForwardDiff(), AutoZygote(), + AutoReverseDiff(), AutoFiniteDiff(), + ) + end + @testset for autodiff in autodiffs_oop @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), StrongWolfe(), @@ -106,9 +117,12 @@ 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(), - ) + autodiffs_iip = if ENZYME_COMPAT + (AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff()) + else + (AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff()) + end + @testset for autodiff in autodiffs_iip @testset "method: $(nameof(typeof(method)))" for method in ( LineSearches.BackTracking(; order = 3), StrongWolfe(), @@ -130,6 +144,9 @@ end using SciMLBase using ADTypes, Tracker, ForwardDiff, Zygote, Enzyme, ReverseDiff, FiniteDiff + # Enzyme is exempted from Julia v1 (stable) testing + const ENZYME_COMPAT = VERSION >= v"1.11" + @testset "OOP Problem" begin nlf(x, p) = x .^ 2 .- p nlp = NonlinearProblem(nlf, [-1.0, 1.0], [3.0]) @@ -144,10 +161,18 @@ end @test abs.(u) ≈ sqrt.([3.0, 3.0]) atol = 1.0e-1 end - @testset for autodiff in ( + autodiffs_oop = if ENZYME_COMPAT + ( AutoTracker(), AutoForwardDiff(), AutoZygote(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), ) + else + ( + AutoTracker(), AutoForwardDiff(), AutoZygote(), + AutoReverseDiff(), AutoFiniteDiff(), + ) + end + @testset for autodiff in autodiffs_oop @testset "method: $(nameof(typeof(method)))" for method in ( BackTracking(; order = Val(3), autodiff), BackTracking(; order = Val(2), autodiff), @@ -174,9 +199,12 @@ end @test abs.(u) ≈ sqrt.([3.0, 3.0]) atol = 1.0e-1 end - @testset for autodiff in ( - AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff(), - ) + autodiffs_iip = if ENZYME_COMPAT + (AutoForwardDiff(), AutoEnzyme(), AutoReverseDiff(), AutoFiniteDiff()) + else + (AutoForwardDiff(), AutoReverseDiff(), AutoFiniteDiff()) + end + @testset for autodiff in autodiffs_iip @testset "method: $(nameof(typeof(method)))" for method in ( BackTracking(; order = Val(3), autodiff), BackTracking(; order = Val(2), autodiff),