Add JET.jl static analysis tests

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"]

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

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),