Avoid specializing OperatorEnum

Author: MilesCranmer

src/EvaluateEquationDerivative.jl

@@ -8,7 +8,7 @@ import ..EquationUtilsModule: count_constants, index_constants, NodeIndex
 import ..EvaluateEquationModule: deg0_eval
 
 function assert_autodiff_enabled(operators::OperatorEnum)
-    if operators.diff_binops === nothing && operators.diff_unaops === nothing
+    if length(operators.diff_binops) == 0 && length(operators.diff_unaops) == 0
         error(
             "Found no differential operators. Did you forget to set `enable_autodiff=true` when creating the `OperatorEnum`?",
         )

src/OperatorEnum.jl

@@ -12,12 +12,11 @@ Defines an enum over operators, along with their derivatives.
 - `diff_binops`: A tuple of Zygote-computed derivatives of the binary operators.
 - `diff_unaops`: A tuple of Zygote-computed derivatives of the unary operators.
 """
-struct OperatorEnum{A<:Tuple,B<:Tuple,dA<:Union{Tuple,Nothing},dB<:Union{Tuple,Nothing}} <:
-       AbstractOperatorEnum
-    binops::A
-    unaops::B
-    diff_binops::dA
-    diff_unaops::dB
+struct OperatorEnum <: AbstractOperatorEnum
+    binops::Vector{Function}
+    unaops::Vector{Function}
+    diff_binops::Vector{Function}
+    diff_unaops::Vector{Function}
 end
 
 """
@@ -30,9 +29,9 @@ Defines an enum over operators, along with their derivatives.
 - `diff_binops`: A tuple of Zygote-computed derivatives of the binary operators.
 - `diff_unaops`: A tuple of Zygote-computed derivatives of the unary operators.
 """
-struct GenericOperatorEnum{A<:Tuple,B<:Tuple} <: AbstractOperatorEnum
-    binops::A
-    unaops::B
+struct GenericOperatorEnum <: AbstractOperatorEnum
+    binops::Vector{Function}
+    unaops::Vector{Function}
 end
 
 end

src/OperatorEnumConstruction.jl

@@ -217,12 +217,13 @@ function OperatorEnum(;
     define_helper_functions::Bool=true,
 )
     @assert length(binary_operators) > 0 || length(unary_operators) > 0
-    binary_operators = Tuple(binary_operators)
-    unary_operators = Tuple(unary_operators)
+
+    binary_operators = convert(Vector{Function}, collect(binary_operators))
+    unary_operators = convert(Vector{Function}, collect(unary_operators))
 
     if enable_autodiff
-        diff_binary_operators = Any[]
-        diff_unary_operators = Any[]
+        diff_binary_operators = Function[]
+        diff_unary_operators = Function[]
 
         test_inputs = map(x -> convert(Float32, x), LinRange(-100, 100, 99))
         # Create grid over [-100, 100]^2:
@@ -259,13 +260,11 @@ function OperatorEnum(;
                 break
             end
         end
-        diff_binary_operators = Tuple(diff_binary_operators)
-        diff_unary_operators = Tuple(diff_unary_operators)
     end
 
     if !enable_autodiff
-        diff_binary_operators = nothing
-        diff_unary_operators = nothing
+        diff_binary_operators = Function[]
+        diff_unary_operators = Function[]
     end
 
     operators = OperatorEnum(
@@ -300,9 +299,6 @@ and `(::Node)(X)`.
 function GenericOperatorEnum(;
     binary_operators=[], unary_operators=[], define_helper_functions::Bool=true
 )
-    binary_operators = Tuple(binary_operators)
-    unary_operators = Tuple(unary_operators)
-
     @assert length(binary_operators) > 0 || length(unary_operators) > 0
 
     operators = GenericOperatorEnum(binary_operators, unary_operators)