Start implementing trimming test

We introduce `SafeTestsets` as part of this, inspired by the way the
downstream tests are set up in the `SciMLBase` repo.

Author: RomeoV

Project.toml

@@ -148,6 +148,7 @@ ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SciMLLogging = "a6db7da4-7206-11f0-1eab-35f2a5dbe1d1"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
@@ -159,4 +160,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "BandedMatrices", "BenchmarkTools", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SparseMatrixColorings", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote", "ReverseDiff", "Tracker", "SciMLLogging"]
+test = ["Aqua", "BandedMatrices", "BenchmarkTools", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SafeTestsets", "SIAMFANLEquations", "SparseConnectivityTracer", "SparseMatrixColorings", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote", "ReverseDiff", "Tracker", "SciMLLogging"]

test/runtests.jl

@@ -1,4 +1,6 @@
-using ReTestItems, NonlinearSolve, Hwloc, InteractiveUtils, Pkg
+using NonlinearSolve, Hwloc, InteractiveUtils, Pkg
+using SafeTestsets
+using ReTestItems
 
 @info sprint(InteractiveUtils.versioninfo)
 
@@ -28,6 +30,13 @@ const GROUP = lowercase(get_from_test_args_or_env("GROUP", "all"))
 # To re-enable: change condition to `true` or `VERSION < v"1.13"`
 const ENZYME_ENABLED = VERSION < v"1.12"
 
+function activate_trim_env!()
+    Pkg.activate(abspath(joinpath(dirname(@__FILE__), "trim")))
+    Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
+    Pkg.instantiate()
+    return nothing
+end
+
 const EXTRA_PKGS = Pkg.PackageSpec[]
 if GROUP == "all" || GROUP == "downstream"
     push!(EXTRA_PKGS, Pkg.PackageSpec("ModelingToolkit"))
@@ -68,8 +77,22 @@ end
 
 @info "Running tests for group: $(GROUP) with $(RETESTITEMS_NWORKERS) workers"
 
-ReTestItems.runtests(
-    NonlinearSolve; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]),
-    nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS,
-    testitem_timeout = 3600
-)
+if GROUP != "trim"
+    ReTestItems.runtests(
+        NonlinearSolve; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]),
+        nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS,
+        testitem_timeout = 3600
+    )
+elseif GROUP == "trim" && VERSION >= v"1.12.0-rc1"  # trimming has been introduced in julia 1.12
+    activate_trim_env!()
+    @safetestset "Clean implementation (non-trimmable)" begin
+        using SciMLBase: successful_retcode
+        include("trim/clean_optimization.jl")
+        @test successful_retcode(minimize(1.0).retcode)
+    end
+    @safetestset "Trimmable implementation" begin
+        using SciMLBase: successful_retcode
+        include("trim/trimmable_optimization.jl")
+        @test successful_retcode(minimize(1.0).retcode)
+    end
+end

test/trim/Project.toml

@@ -0,0 +1,27 @@
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+NonlinearSolveFirstOrder = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d"
+Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
+PolyesterWeave = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+
+[sources]
+ForwardDiff = {rev = "rv/remove-quote-assert-string-interpolation", url = "https://github.com/RomeoV/ForwardDiff.jl"}
+LinearSolve = {rev = "rv/remove-linsolve-forwarddiff-special-path", url = "https://github.com/RomeoV/LinearSolve.jl"}
+NonlinearSolveFirstOrder = {path = "../../lib/NonlinearSolveFirstOrder"}
+Polyester = {rev = "master", url = "https://github.com/RomeoV/Polyester.jl"}
+PolyesterWeave = {rev = "main", url = "https://github.com/RomeoV/PolyesterWeave.jl"}
+SciMLBase = {rev = "as/fix-jet-opt", url = "https://github.com/AayushSabharwal/SciMLBase.jl"}
+
+[compat]
+ADTypes = "1.15.0"
+DiffEqBase = "6.179.0"
+LinearAlgebra = "1.12.0"
+NonlinearSolveFirstOrder = "1.6.0"
+StaticArrays = "1.9.0"

test/trim/clean_optimization.jl

@@ -0,0 +1,32 @@
+using NonlinearSolveFirstOrder
+using ADTypes: AutoForwardDiff
+using ForwardDiff
+using LinearAlgebra
+using StaticArrays
+using LinearSolve
+const LS = LinearSolve
+
+function f(u, p)
+    L, U = cholesky(p.Σ)
+    rhs = (u .* u .- p.λ)
+    # there are some issues currently with LinearSolve and triangular matrices,
+    # so we just make `L` dense here.
+    linprob = LinearProblem(Matrix(L), rhs)
+    alg = LS.GenericLUFactorization()
+    sol = LinearSolve.solve(linprob, alg)
+    return sol.u
+end
+
+struct MyParams{T, M}
+    λ::T
+    Σ::M
+end
+
+function minimize(x)
+    autodiff = AutoForwardDiff(; chunksize=1)
+    alg = TrustRegion(; autodiff, linsolve=LS.CholeskyFactorization())
+    ps = MyParams(rand(), hermitianpart(rand(2,2)+2I))
+    prob = NonlinearLeastSquaresProblem{false}(f, rand(2), ps)
+    sol = solve(prob, alg)
+    return sol
+end

test/trim/trimmable_optimization.jl

@@ -0,0 +1,35 @@
+using NonlinearSolveFirstOrder
+using ADTypes: AutoForwardDiff
+using ForwardDiff
+using LinearAlgebra
+using StaticArrays
+using LinearSolve
+const LS = LinearSolve
+
+function f(u, p)
+    L, U = cholesky(p.Σ)
+    rhs = (u .* u .- p.λ)
+    # there are some issues currently with LinearSolve and triangular matrices,
+    # so we just make `L` dense here.
+    linprob = LinearProblem(Matrix(L), rhs)
+    alg = LS.GenericLUFactorization()
+    sol = LinearSolve.solve(linprob, alg)
+    return sol.u
+end
+
+struct MyParams{T, M}
+    λ::T
+    Σ::M
+end
+
+const autodiff = AutoForwardDiff(; chunksize = 1)
+const alg = TrustRegion(; autodiff, linsolve = LS.CholeskyFactorization())
+const prob = NonlinearLeastSquaresProblem{false}(f, rand(2), MyParams(rand(), hermitianpart(rand(2, 2) + 2I)))
+const cache = init(prob, alg)
+
+function minimize(x)
+    ps = MyParams(x, hermitianpart(rand(2, 2) + 2I))
+    reinit!(cache, rand(2); p = ps)
+    solve!(cache)
+    return cache
+end