deps: add Optim v2 support (#159)

* Compat: support Optim v2

* fix: wrap InplaceObjective hvp/fghvp/fjvp argument order

* fix: explicit InplaceObjective wrapping

* fix: support Optim v1 InplaceObjective

* test: cover objective arity error; @static InplaceObjective layout

* style: run JuliaFormatter

* Refactor InplaceObjective wrapping

* test: cover Optim extension error path

* fix: drop Optim <1; add NLSolversBase compat + Optim v1 CI

* style: run JuliaFormatter

* fix: support Optim v2 (NLSolversBase v8) + add CI smoketest

* style: format Optim extension

* ci: simplify Optim CI and avoid Optim-internal only_fg!

* fix(tests): filter testitems by SR_TEST in runtests

* style: format test runtests

* fix: initialize variable-node val to avoid NaNs

* Revert "fix: initialize variable-node val to avoid NaNs"

This reverts commit 10b59c2.

* fix(tests): include test_optim.jl in "main" test group

The runtests.jl refactor that added SR_TEST filtering
accidentally excluded test_optim.jl from the "main" group.
This meant the regular CI (which resolves Optim v2) ran zero
Optim extension tests — only the dedicated v1 smoketest
exercised the extension.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: MilesCranmerBot <miles.cranmer.bot@gmail.com>
Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: 4 people

.github/workflows/CI.yml

@@ -99,12 +99,43 @@ jobs:
           path: lcov.info
 
 
+  optim_v1_smoketest:
+    name: Optim v1 (NLSolversBase v7) - ubuntu-latest
+    runs-on: ubuntu-latest
+    timeout-minutes: 60
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: '1'
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - name: Pin Optim v1 + NLSolversBase v7
+        run: |
+          julia --color=yes -e 'import Pkg; Pkg.add("Coverage")'
+          julia --color=yes -e 'import Pkg; Pkg.activate("."); Pkg.add(Pkg.PackageSpec(name="Optim", version="1")); Pkg.add(Pkg.PackageSpec(name="NLSolversBase", version="7")); Pkg.status(["Optim", "NLSolversBase"])'
+        shell: bash
+      - name: Run Optim tests (with coverage)
+        id: run-tests
+        run: |
+          SR_TEST=optim julia --color=yes --threads=auto --check-bounds=yes --depwarn=yes --code-coverage=user -e 'import Coverage; import Pkg; Pkg.activate("."); Pkg.test(coverage=true)'
+          julia --color=yes coverage.jl
+        shell: bash
+      - name: Upload coverage artifact
+        if: steps.run-tests.outcome == 'success'
+        uses: actions/upload-artifact@v6
+        with:
+          name: coverage-optim-v1-${{ runner.os }}-julia-1
+          path: lcov.info
+
+
   codecov:
     name: Upload combined coverage to Codecov
     runs-on: ubuntu-latest
     needs:
       - test
       - additional_tests
+      - optim_v1_smoketest
     steps:
       # Codecov uploader expects a git checkout (commit metadata + repo root)
       - uses: actions/checkout@v4

Project.toml

@@ -18,13 +18,14 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+NLSolversBase = "d41bc354-129a-5804-8e4c-c37616107c6c"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 DynamicExpressionsBumperExt = "Bumper"
 DynamicExpressionsLoopVectorizationExt = "LoopVectorization"
-DynamicExpressionsOptimExt = "Optim"
+DynamicExpressionsOptimExt = ["Optim", "NLSolversBase"]
 DynamicExpressionsSymbolicUtilsExt = "SymbolicUtils"
 DynamicExpressionsZygoteExt = "Zygote"
 
@@ -36,7 +37,8 @@ DispatchDoctor = "0.4"
 Interfaces = "0.3"
 LoopVectorization = "0.12"
 MacroTools = "0.4, 0.5"
-Optim = "0.19, 1"
+Optim = "1, 2"
+NLSolversBase = "7, 8"
 PrecompileTools = "1"
 Reexport = "1"
 SymbolicUtils = "4"
@@ -49,5 +51,6 @@ TOML = "1"
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+NLSolversBase = "d41bc354-129a-5804-8e4c-c37616107c6c"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

ext/DynamicExpressionsOptimExt.jl

@@ -9,7 +9,8 @@ using DynamicExpressions:
     set_scalar_constants!,
     get_number_type
 
-import Optim: Optim, OptimizationResults, NLSolversBase
+import Optim: Optim, OptimizationResults
+using NLSolversBase: NLSolversBase
 
 #! format: off
 """
@@ -38,41 +39,136 @@ function Optim.minimizer(r::ExpressionOptimizationResults)
 end
 
 """Wrap function or objective with insertion of values of the constant nodes."""
-function wrap_func(
+@inline function _wrap_objective_x_last(
+    ::Nothing, tree::N, refs
+) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
+    return nothing
+end
+@inline function _wrap_objective_x_last(
     f::F, tree::N, refs
 ) where {F<:Function,T,N<:Union{AbstractExpressionNode{T},AbstractExpression{T}}}
     function wrapped_f(args::Vararg{Any,M}) where {M}
-        first_args = args[begin:(end - 1)]
-        x = args[end]
+        x = args[M]
         set_scalar_constants!(tree, x, refs)
-        return @inline(f(first_args..., tree))
+        newargs = Base.setindex(args, tree, M)
+        return @inline(f(newargs...))
     end
-    # without first args, it looks like this
-    # function wrapped_f(x)
-    #     set_scalar_constants!(tree, x, refs)
-    #     return @inline(f(tree))
-    # end
     return wrapped_f
 end
+
+@inline function _wrap_objective_xv_tail(
+    ::Nothing, tree::N, refs
+) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
+    return nothing
+end
+@inline function _wrap_objective_xv_tail(
+    f::F, tree::N, refs
+) where {F<:Function,T,N<:Union{AbstractExpressionNode{T},AbstractExpression{T}}}
+    function wrapped_f(args::Vararg{Any,M}) where {M}
+        if M < 2
+            throw(
+                ArgumentError(
+                    "Expected at least 2 arguments for objective functions of the form (..., x, v).",
+                ),
+            )
+        end
+        x = args[M - 1]
+        set_scalar_constants!(tree, x, refs)
+        newargs = Base.setindex(args, tree, M - 1)
+        return @inline(f(newargs...))
+    end
+    return wrapped_f
+end
+
+function wrap_func(
+    f::F, tree::N, refs
+) where {F<:Function,T,N<:Union{AbstractExpressionNode{T},AbstractExpression{T}}}
+    return _wrap_objective_x_last(f, tree, refs)
+end
 function wrap_func(
     ::Nothing, tree::N, refs
 ) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
     return nothing
 end
+
+# `NLSolversBase.InplaceObjective` is an internal type whose field layout changed
+# between NLSolversBase versions (and therefore between Optim majors).
+#
+# This extension supports:
+# - Optim v1.x (NLSolversBase v7.x): df, fdf, fgh, hv, fghv
+# - Optim v2.x (NLSolversBase v8.x): fdf, fgh, hvp, fghvp, fjvp
+#
+# We store the fields both as symbols (for runtime layout checks) and as `Val`s
+# (so the wrapper construction is type-stable and can compile-in the field set).
+const _INPLACEOBJECTIVE_SPEC_V8 = (
+    field_syms=(:fdf, :fgh, :hvp, :fghvp, :fjvp),
+    fields=(Val(:fdf), Val(:fgh), Val(:hvp), Val(:fghvp), Val(:fjvp)),
+    x_last=(Val(:fdf), Val(:fgh)),
+    xv_tail=(Val(:hvp), Val(:fghvp), Val(:fjvp)),
+)
+const _INPLACEOBJECTIVE_SPEC_V7 = (
+    field_syms=(:df, :fdf, :fgh, :hv, :fghv),
+    fields=(Val(:df), Val(:fdf), Val(:fgh), Val(:hv), Val(:fghv)),
+    x_last=(Val(:df), Val(:fdf), Val(:fgh)),
+    xv_tail=(Val(:hv), Val(:fghv)),
+)
+
+@inline function _wrap_inplaceobjective_field(
+    v_field::Val{field}, f::NLSolversBase.InplaceObjective, tree::N, refs, spec
+) where {field,N<:Union{AbstractExpressionNode,AbstractExpression}}
+    if v_field in spec.x_last
+        return _wrap_objective_x_last(getfield(f, field), tree, refs)
+    elseif v_field in spec.xv_tail
+        return _wrap_objective_xv_tail(getfield(f, field), tree, refs)
+    else
+        throw(
+            ArgumentError(
+                "Internal error: no wrapping rule for InplaceObjective field $(field). " *
+                "Please open an issue at github.com/SymbolicML/DynamicExpressions.jl with your versions.",
+            ),
+        )
+    end
+end
+
+@inline function _wrap_inplaceobjective(
+    f::NLSolversBase.InplaceObjective, tree::N, refs, spec
+) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
+    wrapped = map(spec.fields) do v_field
+        _wrap_inplaceobjective_field(v_field, f, tree, refs, spec)
+    end
+    return NLSolversBase.InplaceObjective(wrapped...)
+end
+
 function wrap_func(
     f::NLSolversBase.InplaceObjective, tree::N, refs
 ) where {N<:Union{AbstractExpressionNode,AbstractExpression}}
-    # Some objectives, like `Optim.only_fg!(fg!)`, are not functions but instead
+    # Some objectives, like `only_fg!(fg!)`, are not functions but instead
     # `InplaceObjective`. These contain multiple functions, each of which needs to be
     # wrapped. Some functions are `nothing`; those can be left as-is.
-    @assert fieldnames(NLSolversBase.InplaceObjective) == (:df, :fdf, :fgh, :hv, :fghv)
-    return NLSolversBase.InplaceObjective(
-        wrap_func(f.df, tree, refs),
-        wrap_func(f.fdf, tree, refs),
-        wrap_func(f.fgh, tree, refs),
-        wrap_func(f.hv, tree, refs),
-        wrap_func(f.fghv, tree, refs),
-    )
+    #
+    # We use `@static` branching so that only the relevant layout for the *installed*
+    # NLSolversBase version is compiled/instrumented.
+    @static if fieldnames(NLSolversBase.InplaceObjective) ==
+        _INPLACEOBJECTIVE_SPEC_V8.field_syms
+        # NLSolversBase v8 / Optim v2
+        return _wrap_inplaceobjective(f, tree, refs, _INPLACEOBJECTIVE_SPEC_V8)
+    elseif fieldnames(NLSolversBase.InplaceObjective) ==
+        _INPLACEOBJECTIVE_SPEC_V7.field_syms
+        # NLSolversBase v7 / Optim v1
+        return _wrap_inplaceobjective(f, tree, refs, _INPLACEOBJECTIVE_SPEC_V7)
+        # (Optim < 1 is no longer supported.)
+    else
+        # LCOV_EXCL_START
+        fields = fieldnames(NLSolversBase.InplaceObjective)
+        throw(
+            ArgumentError(
+                "Unsupported NLSolversBase.InplaceObjective field layout: $(fields). " *
+                "This extension supports layouts used by NLSolversBase v7 (Optim v1) and v8 (Optim v2). " *
+                "Please open an issue at github.com/SymbolicML/DynamicExpressions.jl with your versions.",
+            ),
+        )
+        # LCOV_EXCL_END
+    end
 end
 
 """

test/Project.toml

@@ -11,6 +11,7 @@ Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+NLSolversBase = "d41bc354-129a-5804-8e4c-c37616107c6c"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"

test/runtests.jl

@@ -1,21 +1,32 @@
 using SafeTestsets
 using TestItemRunner
 
-# Check if SR_ENZYME_TEST is set in env
-test_name = split(get(ENV, "SR_TEST", "main"), ",")
+# Control which test groups run.
+#
+# Accepts a comma-separated list in SR_TEST (default: "main").
+#
+# - "main": full test suite (testitems)
+# - "optim": Optim-specific testitems only
+# - "narity": n-ary testitems only
+# - "jet": JET analysis
+# - "enzyme": Enzyme tests
 
-unknown_tests = filter(Base.Fix2(∉, ["enzyme", "jet", "main", "narity"]), test_name)
+test_names = split(get(ENV, "SR_TEST", "main"), ",")
+
+allowed = ["enzyme", "jet", "main", "narity", "optim"]
+unknown_tests = filter(Base.Fix2(∉, allowed), test_names)
 
 if !isempty(unknown_tests)
     error("Unknown test names: $unknown_tests")
 end
 
-if "enzyme" in test_name
+if "enzyme" in test_names
     @safetestset "Test enzyme derivatives" begin
         include("test_enzyme.jl")
     end
 end
-if "jet" in test_name
+
+if "jet" in test_names
     @safetestset "JET" begin
         using Preferences
         set_preferences!(
@@ -61,7 +72,25 @@ if "jet" in test_name
         end
     end
 end
-if "main" in test_name
-    include("unittest.jl")
-    @run_package_tests
+
+# TestItemRunner's `@run_package_tests` scans *all* `.jl` files under the package root,
+# so we must filter to only the testitem files we actually want to run.
+# (Simply `include(...)`-ing a subset of files is not sufficient.)
+
+testitem_suffixes = String[]
+
+if "main" in test_names
+    push!(testitem_suffixes, joinpath("test", "unittest.jl"))
+    push!(testitem_suffixes, joinpath("test", "test_optim.jl"))
+end
+if "optim" in test_names
+    push!(testitem_suffixes, joinpath("test", "test_optim.jl"))
+end
+if "narity" in test_names
+    push!(testitem_suffixes, joinpath("test", "test_n_arity_nodes.jl"))
+end
+
+if !isempty(testitem_suffixes)
+    @run_package_tests filter =
+        ti -> any(suf -> endswith(ti.filename, suf), testitem_suffixes)
 end