diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index f11cae89..de4709e5 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -97,8 +97,8 @@ jobs: path: lcov.info - optim_v1_smoketest: - name: Optim v1 (NLSolversBase v7) - ubuntu-latest + optim_lower_bound_smoketest: + name: Optim lower bound (Optim v2.0.0, NLSolversBase v8.0.0) - ubuntu-latest runs-on: ubuntu-latest timeout-minutes: 60 steps: @@ -108,10 +108,10 @@ jobs: version: '1' - uses: julia-actions/cache@v2 - uses: julia-actions/julia-buildpkg@v1 - - name: Pin Optim v1 + NLSolversBase v7 + - name: Pin Optim lower bounds 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"])' + julia --color=yes -e 'import Pkg; Pkg.activate("."); Pkg.add(Pkg.PackageSpec(name="Optim", version="2.0.0")); Pkg.add(Pkg.PackageSpec(name="NLSolversBase", version="8.0.0")); Pkg.status(["Optim", "NLSolversBase"])' shell: bash - name: Run Optim tests (with coverage) id: run-tests @@ -123,7 +123,7 @@ jobs: if: steps.run-tests.outcome == 'success' uses: actions/upload-artifact@v6 with: - name: coverage-optim-v1-${{ runner.os }}-julia-1 + name: coverage-optim-lower-bound-${{ runner.os }}-julia-1 path: lcov.info @@ -133,7 +133,7 @@ jobs: needs: - test - additional_tests - - optim_v1_smoketest + - optim_lower_bound_smoketest steps: # Codecov uploader expects a git checkout (commit metadata + repo root) - uses: actions/checkout@v4 diff --git a/Project.toml b/Project.toml index 0da0de6a..1df6bc8e 100644 --- a/Project.toml +++ b/Project.toml @@ -35,13 +35,13 @@ ChainRulesCore = "1.25.1" Compat = "4.16" DispatchDoctor = "0.4" Interfaces = "0.3" -LoopVectorization = "0.12" +LoopVectorization = "0.12.172" MacroTools = "0.5.16" -Optim = "1, 2" -NLSolversBase = "7, 8" +Optim = "2" +NLSolversBase = "8" PrecompileTools = "1.2.1" Reexport = "1.2.2" -SymbolicUtils = "4" +SymbolicUtils = "4.1" Zygote = "0.7" julia = "1.10" Random = "1" diff --git a/src/DynamicExpressions.jl b/src/DynamicExpressions.jl index 355e7b98..d4206051 100644 --- a/src/DynamicExpressions.jl +++ b/src/DynamicExpressions.jl @@ -73,7 +73,8 @@ import .NodeModule: has_constants, count_scalar_constants, get_scalar_constants, - set_scalar_constants! + set_scalar_constants!, + set_scalar_constants @reexport import .StringsModule: string_tree, print_tree import .StringsModule: get_op_name, get_pretty_op_name @reexport import .OperatorEnumModule: AbstractOperatorEnum diff --git a/src/Expression.jl b/src/Expression.jl index e434b206..e288e2b0 100644 --- a/src/Expression.jl +++ b/src/Expression.jl @@ -19,7 +19,8 @@ import ..NodeUtilsModule: has_constants, count_scalar_constants, get_scalar_constants, - set_scalar_constants! + set_scalar_constants!, + set_scalar_constants import ..NodePreallocationModule: copy_into!, allocate_container import ..EvaluateModule: eval_tree_array, differentiable_eval_tree_array import ..EvaluateDerivativeModule: eval_grad_tree_array @@ -327,6 +328,9 @@ end function set_scalar_constants!(ex::Expression{T}, constants, refs) where {T} return set_scalar_constants!(get_tree(ex), constants, refs) end +function set_scalar_constants(ex::Expression, constants) + return Expression(set_scalar_constants(get_tree(ex), constants), get_metadata(ex)) +end function extract_gradient( gradient::@NamedTuple{tree::NT, metadata::Nothing}, ex::Expression{T,N} ) where {T,N<:AbstractExpressionNode{T},NT<:NodeTangent{T,N}} diff --git a/src/NodeUtils.jl b/src/NodeUtils.jl index 5b75f7b1..4bff3262 100644 --- a/src/NodeUtils.jl +++ b/src/NodeUtils.jl @@ -8,6 +8,7 @@ import ..NodeModule: Node, preserve_sharing, constructorof, + with_type_parameters, set_children!, copy_node, count_nodes, @@ -142,6 +143,29 @@ function set_scalar_constants!(tree::AbstractExpressionNode{T}, constants, refs) return tree end +""" + set_scalar_constants(tree::AbstractExpressionNode{T}, constants) where {T} + +Return a *new* tree with scalar constants set (non-mutating). + +This is equivalent to `copy(tree)` followed by [`set_scalar_constants!`](@ref), +but will also promote the tree's number type to accommodate `eltype(constants)`. +This makes it possible to forward-mode differentiate through constant-setting +(e.g. with `ForwardDiff.Dual` constants). +""" +function set_scalar_constants(tree::AbstractExpressionNode{T}, constants) where {T} + Tc = eltype(constants) + Tout = promote_type(T, Tc) + newtree = if Tout === T + copy(tree) + else + convert(with_type_parameters(typeof(tree), Tout), tree) + end + _, refs = get_scalar_constants(newtree) + set_scalar_constants!(newtree, constants, refs) + return newtree +end + ## Assign index to nodes of a tree # This will mirror a Node struct, rather # than adding a new attribute to Node. diff --git a/test/Project.toml b/test/Project.toml index e4b346d5..da708df8 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -26,7 +26,7 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] -Aqua = "0.8" +Aqua = "0.8.14" JET = "0.9, 0.10" TestItems = "1" TestItemRunner = "1" diff --git a/test/test_utils.jl b/test/test_utils.jl index 7a862cb7..e161e3a5 100644 --- a/test/test_utils.jl +++ b/test/test_utils.jl @@ -1,5 +1,6 @@ using DynamicExpressions using DynamicExpressions.UtilsModule: fill_similar +using ForwardDiff using Test operators = OperatorEnum(2 => (+, *, -, /, ^), 1 => (exp, sin)) @@ -36,6 +37,35 @@ tree = x1 + Node(; val=0.0) - sin(x2 - Node(; val=0.5)) set_scalar_constants!(tree, [1.0, 2.0], [Ref(tree.l.r), Ref(tree.r.l.r)]) @test repr(tree) == "(x1 + 1.0) - sin(x2 - 2.0)" +# Non-mutating set constants (and ForwardDiff friendliness): +let + x1 = Node("x1") + tree = x1 * Node(; val=0.0) + Node(; val=0.0) + @test get_scalar_constants(tree)[1] == [0.0, 0.0] + X = reshape([1.0, 2.0, 3.0], 1, :) + operators = OperatorEnum(2 => (+, *), 1 => (sin,)) + + f(c) = begin + t2 = set_scalar_constants(tree, c) + return sum(eval_tree_array(t2, X, operators)[1]) + end + + g = ForwardDiff.gradient(f, [2.0, 3.0]) + @test g ≈ [6.0, 3.0] + + # Original tree unchanged. + @test get_scalar_constants(tree)[1] == [0.0, 0.0] + + # Expression wrapper also works (including promotion): + ex = Expression(tree; operators=operators, variable_names=["x1"]) + f_ex(c) = begin + ex2 = set_scalar_constants(ex, c) + return sum(eval_tree_array(get_tree(ex2), X, operators)[1]) + end + g_ex = ForwardDiff.gradient(f_ex, [2.0, 3.0]) + @test g_ex ≈ [6.0, 3.0] +end + # Ensure that fill_similar is type stable x = randn(Float32, 3, 10) @inferred fill_similar(0.5f0, x, axes(x, 1))