Fix inference test

Author: MilesCranmer

src/EvaluateEquationDerivative.jl

@@ -28,7 +28,7 @@ respect to `x1`.
 - `cX::AbstractMatrix{T}`: The data matrix, with each column being a data point.
 - `operators::OperatorEnum`: The operators used to create the `tree`.
 - `direction::Integer`: The index of the variable to take the derivative with respect to.
-- `turbo::Bool`: Use `LoopVectorization.@turbo` for faster evaluation.
+- `turbo::Union{Val,Bool}`: Use `LoopVectorization.@turbo` for faster evaluation.
 
 # Returns
 
@@ -40,7 +40,7 @@ function eval_diff_tree_array(
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
     direction::Integer;
-    turbo::Bool=false,
+    turbo::Union{Val,Bool}=Val(false),
 ) where {T<:Number}
     # TODO: Implement quick check for whether the variable is actually used
     # in this tree. Otherwise, return zero.

src/EvaluationHelpers.jl

@@ -59,11 +59,9 @@ end
 
 # Gradients:
 function _grad_evaluator(
-    tree::AbstractExpressionNode, X, operators::OperatorEnum; variable=true, kws...
+    tree::AbstractExpressionNode, X, operators::OperatorEnum; variable=Val(true), kws...
 )
-    _, grad, did_complete = eval_grad_tree_array(
-        tree, X, operators; variable=variable, kws...
-    )
+    _, grad, did_complete = eval_grad_tree_array(tree, X, operators; variable, kws...)
     !did_complete && (grad .= convert(eltype(grad), NaN))
     return grad
 end
@@ -74,7 +72,7 @@ function _grad_evaluator(
 end
 
 """
-    (tree::AbstractExpressionNode{T})'(X::AbstractMatrix{T}, operators::OperatorEnum; turbo::Bool=false, variable::Bool=true)
+    (tree::AbstractExpressionNode{T})'(X::AbstractMatrix{T}, operators::OperatorEnum; turbo::Union{Bool,Val}=Val(false), variable::Union{Bool,Val}=Val(true))
 
 Compute the forward-mode derivative of an expression, using a similar
 structure and optimization to eval_tree_array. `variable` specifies whether
@@ -84,9 +82,9 @@ to every constant in the expression.
 # Arguments
 - `X::AbstractMatrix{T}`: The data matrix, with each column being a data point.
 - `operators::OperatorEnum`: The operators used to create the `tree`.
-- `variable::Bool`: Whether to take derivatives with respect to features (i.e., `X` - with `variable=true`),
+- `variable::Union{Bool,Val}`: Whether to take derivatives with respect to features (i.e., `X` - with `variable=true`),
     or with respect to every constant in the expression (`variable=false`).
-- `turbo::Bool`: Use `LoopVectorization.@turbo` for faster evaluation.
+- `turbo::Union{Bool,Val}`: Use `LoopVectorization.@turbo` for faster evaluation.
 
 # Returns
 

test/test_derivatives.jl

@@ -37,8 +37,8 @@ function array_test(ar1, ar2; rtol=0.1)
     return isapprox(ar1, ar2; rtol=rtol)
 end
 
-for type in [Float16, Float32, Float64], turbo in [true, false]
-    type == Float16 && turbo && continue
+for type in [Float16, Float32, Float64], turbo in [Val(true), Val(false)]
+    type == Float16 && turbo isa Val{true} && continue
 
     println(
         "Testing derivatives with respect to variables, with type=$(type) and turbo=$(turbo).",
@@ -144,15 +144,15 @@ for type in [Float16, Float32, Float64], turbo in [true, false]
 end
 
 @testset "NodeIndex" begin
-    import DynamicExpressions: get_constants, NodeIndex, index_constants
+    @eval import DynamicExpressions: get_constants, NodeIndex, index_constants
 
     operators = OperatorEnum(;
         binary_operators=(+, *, -, /, pow_abs2), unary_operators=(custom_cos, exp, sin)
     )
     @extend_operators operators
     tree = equation3(nx1, nx2, nx3)
 
-    """Check whether the ordering of constant_list is the same as the ordering of node_index."""
+    # Check whether the ordering of constant_list is the same as the ordering of node_index.
     @eval function check_tree(
         tree::Node, node_index::NodeIndex, constant_list::AbstractVector
     )

test/test_evaluation.jl

@@ -36,9 +36,10 @@ functions = [
     (x1, x2, x3) -> (sin(cos(sin(cos(x1) * x3) * 3.0) * -0.5) + 2.0) * 5.0,
 ]
 
-for turbo in [false, true], T in [Float16, Float32, Float64, ComplexF32, ComplexF64]
+for turbo in [Val(false), Val(true)],
+    T in [Float16, Float32, Float64, ComplexF32, ComplexF64]
     # Float16 not implemented:
-    turbo && !(T in (Float32, Float64)) && continue
+    turbo isa Val{true} && !(T in (Float32, Float64)) && continue
     @testset "Test evaluation of trees with turbo=$turbo, T=$T" begin
         for (i_func, fnc) in enumerate(functions)