Expose recorder function for DynamicAutodiff.jl

Project.toml

@@ -48,7 +48,7 @@ Dates = "1"
 DifferentiationInterface = "0.5, 0.6"
 DispatchDoctor = "^0.4.17"
 Distributed = "<0.0.1, 1"
-DynamicExpressions = "1.5.0"
+DynamicExpressions = "1.6.0"
 DynamicQuantities = "1"
 Enzyme = "0.12, 0.13"
 JSON3 = "1"

src/SymbolicRegression.jl

@@ -131,6 +131,8 @@ using DynamicExpressions:
     with_contents,
     with_metadata
 using DynamicExpressions: with_type_parameters
+# TODO: Reexport D once DynamicAutodiff is registered
+# @reexport using DynamicAutodiff: D
 @reexport using LossFunctions:
     MarginLoss,
     DistanceLoss,

src/TemplateExpression.jl

@@ -39,6 +39,9 @@ using ..MutateModule: MutateModule as MM
 using ..PopMemberModule: PopMember
 using ..ComposableExpressionModule: ComposableExpression, ValidVector
 
+# TODO: Modify `D` once DynamicAutodiff is registered
+# import DynamicAutodiff: D
+
 """
     TemplateStructure{K,E,NF} <: Function
 
@@ -86,14 +89,26 @@ function _record_composable_expression!(variable_constraints, ::Val{k}, args...)
     return isempty(args) ? 0.0 : first(args)
 end
 
+struct ArgumentRecorder{F} <: Function
+    f::F
+end
+(f::ArgumentRecorder)(args...) = f.f(args...)
+
+# TODO: Modify `D` once DynamicAutodiff is registered
+# We pass through the derivative operators, since
+# we just want to record the number of arguments.
+# DA.D(f::ArgumentRecorder, _::Integer) = f
+
 """Infers number of features used by each subexpression, by passing in test data."""
 function infer_variable_constraints(::Val{K}, combiner::F) where {K,F}
     variable_constraints = NamedTuple{K}(map(_ -> Ref(-1), K))
     # Now, we need to evaluate the `combine` function to see how many
     # features are used for each function call. If unset, we record it.
     # If set, we validate.
     inner = Fix{1}(_record_composable_expression!, variable_constraints)
-    _recorders_of_composable_expressions = NamedTuple{K}(map(k -> Fix{1}(inner, Val(k)), K))
+    _recorders_of_composable_expressions = NamedTuple{K}(
+        map(k -> ArgumentRecorder(Fix{1}(inner, Val(k))), K)
+    )
     # We use an evaluation to get the variable constraints
     combiner(
         _recorders_of_composable_expressions,