Add fast-compilation branches for gradients too

Author: MilesCranmer

src/EvaluateEquationDerivative.jl

@@ -4,7 +4,8 @@ import ..EquationModule: AbstractExpressionNode, constructorof
 import ..OperatorEnumModule: OperatorEnum
 import ..UtilsModule: is_bad_array, fill_similar
 import ..EquationUtilsModule: count_constants, index_constants, NodeIndex
-import ..EvaluateEquationModule: deg0_eval, get_nuna, get_nbin
+import ..EvaluateEquationModule:
+    deg0_eval, get_nuna, get_nbin, OPERATOR_LIMIT_BEFORE_SLOWDOWN
 import ..ExtensionInterfaceModule: _zygote_gradient
 
 struct ResultOk2{A<:AbstractArray,B<:AbstractArray}
@@ -70,26 +71,42 @@ end
 )::ResultOk2 where {T<:Number}
     nuna = get_nuna(operators)
     nbin = get_nbin(operators)
-    quote
-        result = if tree.degree == 0
-            diff_deg0_eval(tree, cX, direction)
-        elseif tree.degree == 1
-            op_idx = tree.op
+    deg1_branch = if nuna > OPERATOR_LIMIT_BEFORE_SLOWDOWN
+        quote
+            diff_deg1_eval(tree, cX, operators.unaops[op_idx], operators, direction)
+        end
+    else
+        quote
             Base.Cartesian.@nif(
                 $nuna,
                 i -> i == op_idx,
                 i ->
                     diff_deg1_eval(tree, cX, operators.unaops[i], operators, direction)
             )
-        else
-            op_idx = tree.op
+        end
+    end
+    deg2_branch = if nbin > OPERATOR_LIMIT_BEFORE_SLOWDOWN
+        diff_deg2_eval(tree, cX, operators.binops[op_idx], operators, direction)
+    else
+        quote
             Base.Cartesian.@nif(
                 $nbin,
                 i -> i == op_idx,
                 i ->
                     diff_deg2_eval(tree, cX, operators.binops[i], operators, direction)
             )
         end
+    end
+    quote
+        result = if tree.degree == 0
+            diff_deg0_eval(tree, cX, direction)
+        elseif tree.degree == 1
+            op_idx = tree.op
+            $deg1_branch
+        else
+            op_idx = tree.op
+            $deg2_branch
+        end
         !result.ok && return result
         return ResultOk2(
             result.x, result.dx, !(is_bad_array(result.x) || is_bad_array(result.dx))
@@ -251,39 +268,57 @@ end
 )::ResultOk2 where {T<:Number,variable}
     nuna = get_nuna(operators)
     nbin = get_nbin(operators)
+    deg1_branch_skeleton = quote
+        grad_deg1_eval(
+            tree,
+            n_gradients,
+            index_tree,
+            cX,
+            operators.unaops[i],
+            operators,
+            Val(variable),
+        )
+    end
+    deg2_branch_skeleton = quote
+        grad_deg2_eval(
+            tree,
+            n_gradients,
+            index_tree,
+            cX,
+            operators.binops[i],
+            operators,
+            Val(variable),
+        )
+    end
+    deg1_branch = if nuna > OPERATOR_LIMIT_BEFORE_SLOWDOWN
+        quote
+            i = tree.op
+            $deg1_branch_skeleton
+        end
+    else
+        quote
+            op_idx = tree.op
+            Base.Cartesian.@nif($nuna, i -> i == op_idx, i -> $deg1_branch_skeleton)
+        end
+    end
+    deg2_branch = if nbin > OPERATOR_LIMIT_BEFORE_SLOWDOWN
+        quote
+            i = tree.op
+            $deg2_branch_skeleton
+        end
+    else
+        quote
+            op_idx = tree.op
+            Base.Cartesian.@nif($nbin, i -> i == op_idx, i -> $deg2_branch_skeleton)
+        end
+    end
     quote
         if tree.degree == 0
             grad_deg0_eval(tree, n_gradients, index_tree, cX, Val(variable))
         elseif tree.degree == 1
-            op_idx = tree.op
-            Base.Cartesian.@nif(
-                $nuna,
-                i -> i == op_idx,
-                i -> grad_deg1_eval(
-                    tree,
-                    n_gradients,
-                    index_tree,
-                    cX,
-                    operators.unaops[i],
-                    operators,
-                    Val(variable),
-                )
-            )
+            $deg1_branch
         else
-            op_idx = tree.op
-            Base.Cartesian.@nif(
-                $nbin,
-                i -> i == op_idx,
-                i -> grad_deg2_eval(
-                    tree,
-                    n_gradients,
-                    index_tree,
-                    cX,
-                    operators.binops[i],
-                    operators,
-                    Val(variable),
-                )
-            )
+            $deg2_branch
         end
     end
 end

test/test_derivatives.jl

@@ -142,28 +142,63 @@ for type in [Float16, Float32, Float64], turbo in [true, false]
     println("Done.")
 end
 
-println("Testing NodeIndex.")
-
-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."""
-function check_tree(tree::Node, node_index::NodeIndex, constant_list::AbstractVector)
-    if tree.degree == 0
-        (!tree.constant) || tree.val == constant_list[node_index.val::UInt16]
-    elseif tree.degree == 1
-        check_tree(tree.l, node_index.l, constant_list)
-    else
-        check_tree(tree.l, node_index.l, constant_list) &&
-            check_tree(tree.r, node_index.r, constant_list)
+@testset "NodeIndex" begin
+    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."""
+    @eval function check_tree(
+        tree::Node, node_index::NodeIndex, constant_list::AbstractVector
+    )
+        if tree.degree == 0
+            (!tree.constant) || tree.val == constant_list[node_index.val::UInt16]
+        elseif tree.degree == 1
+            check_tree(tree.l, node_index.l, constant_list)
+        else
+            check_tree(tree.l, node_index.l, constant_list) &&
+                check_tree(tree.r, node_index.r, constant_list)
+        end
     end
+
+    @test check_tree(tree, index_constants(tree), get_constants(tree))
 end
 
-@test check_tree(tree, index_constants(tree), get_constants(tree))
+@testset "Test many operators" begin
+    # Since we use `@nif` in evaluating expressions,
+    # we can see if there are any issues with LARGE numbers of operators.
+    num_ops = 100
+    binary_operators = [@eval function (x, y)
+        return x + y
+    end for i in 1:num_ops]
+    unary_operators = [@eval function (x)
+        return x^2
+    end for i in 1:num_ops]
+
+    # This OperatorEnum will trigger the fallback code for fast compilation.
+    many_ops_operators = OperatorEnum(;
+        binary_operators=cat([+, -, *, /], binary_operators; dims=1),
+        unary_operators=cat([sin, cos], unary_operators; dims=1),
+    )
 
-println("Done.")
+    # This OperatorEnum will go through the regular evaluation code.
+    only_basic_ops_operator = OperatorEnum(;
+        binary_operators=[+, -, *, /], unary_operators=[sin, cos]
+    )
+
+    # We want to compare their gradients
+    num_tests = 100
+    n_features = 3
+    for _ in 1:num_tests
+        tree = gen_random_tree_fixed_size(20, only_basic_ops_operator, n_features, Float64)
+        X = randn(Float64, n_features, 10)
+        basic_eval = tree'(X, only_basic_ops_operator)
+        many_ops_eval = tree'(X, many_ops_operators)
+        @test (all(isnan, basic_eval) && all(isnan, many_ops_eval)) ||
+            basic_eval ≈ many_ops_eval
+    end
+end