feat: add non-mutating set_scalar_constants

Author: MilesCranmerBot

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

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}}

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.

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