ci: add extra benchmark for parametric nodes

Author: MilesCranmer

benchmark/benchmarks.jl

@@ -213,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
@@ -222,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
 

src/ParametricExpression.jl

@@ -173,26 +173,8 @@ end
 ###############################################################################
 # Extra utilities for parametric-specific behavior ############################
 ###############################################################################
-## As explained in AbstractExpressionNode, we can implement custom behavior for
-## the parametric expression by implementing the following methods:
-# - `count_nodes`
-# - `count_constants`
-# - `count_depth`
-# - `index_constant_nodes`
-# - `has_operators`
-# - `has_constants`
-# - `get_scalar_constants`
-# - `set_scalar_constants!`
-# - `string_tree`
-# - `max_feature`
-# - `eval_tree_array`
-# - `eval_grad_tree_array`
-# - `_grad_evaluator`
-
-## For a parametric struct, we only wish to implement the following
 
 #! format: off
-
 struct InterfaceError <: Exception
 end
 _interface_error() = throw(InterfaceError())