Merge 455e80f into cab1143

Author: gca30

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <miles.cranmer@gmail.com>"]
-version = "0.18.5"
+version = "0.18.6"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -14,7 +14,6 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
-TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 
 [weakdeps]
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
@@ -43,7 +42,6 @@ PackageExtensionCompat = "1"
 PrecompileTools = "1"
 Reexport = "1"
 SymbolicUtils = "0.19, ^1.0.5, 2"
-TestItems = "0.1"
 Zygote = "0.6"
 julia = "1.6"
 

benchmark/benchmarks.jl

@@ -1,10 +1,7 @@
 using DynamicExpressions, BenchmarkTools, Random
 
 # Trigger extensions:
-using LoopVectorization
-using Bumper
-using StrideArrays
-using Zygote
+using LoopVectorization, Bumper, StrideArrays, Zygote
 
 if PACKAGE_VERSION < v"0.14.0"
     @eval using DynamicExpressions: Node as GraphNode
@@ -18,6 +15,14 @@ else
     @eval using DynamicExpressions.NodeUtilsModule: is_constant
 end
 
+if PACKAGE_VERSION < v"0.18.6"
+    @eval using DynamicExpressions:
+        index_constants as index_constant_nodes,
+        count_constants as count_constant_nodes,
+        get_constants as get_scalar_constants,
+        set_constants! as set_scalar_constants!
+end
+
 include("../test/tree_gen_utils.jl")
 
 const SUITE = BenchmarkGroup()
@@ -113,15 +118,16 @@ end
     PACKAGE_VERSION < v"0.14.0" && return :(copy_node(t; preserve_sharing=preserve_sharing))
     return :(copy_node(t))  # Assume type used to infer sharing
 end
-@generated function get_set_constants!(tree::N) where {T,N<:AbstractExpressionNode{T}}
-    if !(@isdefined set_constants!)
-        return :(set_constants(tree, get_constants(tree)))
-    elseif hasmethod(set_constants!, Tuple{N, Vector{T}})
-        return :(set_constants!(tree, get_constants(tree)))
+@generated function get_set_constants!(tree::N) where {N}
+    T = eltype(N)
+    if !(@isdefined set_scalar_constants!)
+        return :(set_scalar_constants(tree, get_scalar_constants(tree)))
+    elseif hasmethod(set_scalar_constants!, Tuple{N, Vector{T}})
+        return :(set_scalar_constants!(tree, get_scalar_constants(tree)))
     else
         return quote
-            let (x, refs) = get_constants(tree)
-                set_constants!(tree, x, refs)
+                let (x, refs) = get_scalar_constants(tree)
+                set_scalar_constants!(tree, x, refs)
             end
         end
     end
@@ -141,12 +147,12 @@ function benchmark_utilities()
         :combine_operators,
         :count_nodes,
         :count_depth,
-        :count_constants,
+        :count_constant_nodes,
         :has_constants,
         :has_operators,
         :is_constant,
         :get_set_constants!,
-        :index_constants,
+        :index_constant_nodes,
         :string_tree,
         :hash,
     )
@@ -157,9 +163,9 @@ function benchmark_utilities()
             [
                 :simplify_tree,
                 :count_nodes,
-                :count_constants,
+                :count_constant_nodes,
                 :get_set_constants!,
-                :index_constants,
+                :index_constant_nodes,
                 :string_tree,
             ],
         )
@@ -207,7 +213,8 @@ function benchmark_utilities()
                     setup=(
                         ntrees=100;
                         n=20;
-                        trees=[$preprocess(gen_random_tree_fixed_size(n, $operators, 5, Float32)) for _ in 1:ntrees]
+                        rng=Random.MersenneTwister(0);
+                        trees=[$preprocess(gen_random_tree_fixed_size(n, $operators, 5, Float32, Node, rng)) for _ in 1:ntrees]
                     )
                 )
                 #! format: on
@@ -216,6 +223,37 @@ function benchmark_utilities()
         end
     end
 
+    # Additional methods
+    @static if PACKAGE_VERSION >= v"0.18.0"
+        suite["get_set_constants_parametric"] = @benchmarkable(
+            [get_set_constants!(ex) for ex in exs],
+            seconds = 10.0,
+            setup = (
+                operators = $operators;
+                ntrees = 100;
+                n = 20;
+                n_features = 5;
+                n_params = 3;
+                n_param_classes = 10;
+                rng = Random.MersenneTwister(0);
+                exs = [
+                    let tree = gen_random_tree_fixed_size(
+                            n, operators, n_features, Float32, ParametricNode, rng
+                        )
+                        ex = ParametricExpression(
+                            tree;
+                            operators,
+                            variable_names=map(i -> "x$i", 1:n_features),
+                            parameters=randn(rng, Float32, n_params, n_param_classes),
+                            parameter_names=map(i -> "p$i", 1:n_params),
+                        )
+                        ex
+                    end for _ in 1:ntrees
+                ]
+            )
+        )
+    end
+
     return suite
 end
 

ext/DynamicExpressionsOptimExt.jl

@@ -5,8 +5,9 @@ using DynamicExpressions:
     AbstractExpressionNode,
     filter_map,
     eval_tree_array,
-    get_constants,
-    set_constants!
+    get_scalar_constants,
+    set_scalar_constants!,
+    get_number_type
 using Compat: @inline
 
 import Optim: Optim, OptimizationResults, NLSolversBase
@@ -44,9 +45,14 @@ function wrap_func(
     function wrapped_f(args::Vararg{Any,M}) where {M}
         first_args = args[begin:(end - 1)]
         x = args[end]
-        set_constants!(tree, x, refs)
+        set_scalar_constants!(tree, x, refs)
         return @inline(f(first_args..., tree))
     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
 function wrap_func(
@@ -100,7 +106,8 @@ function Optim.optimize(
     if make_copy
         tree = copy(tree)
     end
-    x0, refs = get_constants(tree)
+
+    x0, refs = get_scalar_constants(tree)
     if !isnothing(h!)
         throw(
             ArgumentError(
@@ -117,7 +124,7 @@ function Optim.optimize(
         )
     end
     minimizer = Optim.minimizer(base_res)
-    set_constants!(tree, minimizer, refs)
+    set_scalar_constants!(tree, minimizer, refs)
     return ExpressionOptimizationResults(base_res, tree)
 end
 

src/DynamicExpressions.jl

@@ -4,6 +4,7 @@ using DispatchDoctor: @stable, @unstable
 
 @stable default_mode = "disable" begin
     include("Utils.jl")
+    include("ValueInterface.jl")
     include("ExtensionInterface.jl")
     include("OperatorEnum.jl")
     include("Node.jl")
@@ -25,6 +26,13 @@ import PackageExtensionCompat: @require_extensions
 import Reexport: @reexport
 macro ignore(args...) end
 
+import .ValueInterfaceModule:
+    is_valid,
+    is_valid_array,
+    get_number_type,
+    pack_scalar_constants!,
+    unpack_scalar_constants,
+    ValueInterface
 @reexport import .NodeModule:
     AbstractNode,
     AbstractExpressionNode,
@@ -47,14 +55,15 @@ import .NodeModule:
     branch_equal
 @reexport import .NodeUtilsModule:
     count_nodes,
-    count_constants,
+    count_constant_nodes,
     count_depth,
     NodeIndex,
-    index_constants,
+    index_constant_nodes,
     has_operators,
     has_constants,
-    get_constants,
-    set_constants!
+    count_scalar_constants,
+    get_scalar_constants,
+    set_scalar_constants!
 @reexport import .StringsModule: string_tree, print_tree
 @reexport import .OperatorEnumModule: AbstractOperatorEnum
 @reexport import .OperatorEnumConstructionModule: