test: add preliminary JET tests

Author: AayushSabharwal

Project.toml

@@ -100,6 +100,7 @@ julia = "1.10"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
@@ -114,4 +115,4 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "ForwardDiff", "MLStyle", "PartialFunctions", "Pkg", "SafeTestsets", "Serialization", "StableRNGs", "StaticArrays", "Tables", "Test", "UnicodePlots", "Zygote"]
+test = ["Aqua", "ForwardDiff", "MLStyle", "PartialFunctions", "Pkg", "SafeTestsets", "Serialization", "StableRNGs", "StaticArrays", "Tables", "Test", "UnicodePlots", "Zygote", "JET"]

test/JET.jl

@@ -0,0 +1,20 @@
+using NonlinearSolve
+using LinearAlgebra
+using ADTypes
+using JET
+
+function f(u, p)
+    L, U = cholesky(p.Σ)
+    return L \ (u .* u .- p.λ)
+end
+
+function minimize(λ=1.0)
+    ps = (; λ, Σ=hermitianpart(rand(2,2) + 2*I))
+    u₀ = rand(2)
+    prob = NonlinearLeastSquaresProblem{false}(f, u₀, ps)
+    autodiff = AutoForwardDiff(; chunksize=1)
+    sol = solve(prob, SimpleTrustRegion(; autodiff))
+    return sol.u
+end
+
+@test_opt minimize()

test/runtests.jl

@@ -60,6 +60,9 @@ end
         @time @safetestset "Clocks" begin
             include("clock.jl")
         end
+        @time @safetestset "JET" begin
+            include("JET.jl")
+        end
     end
 
     if !is_APPVEYOR &&