feat: avoid double descent in rrule

Author: MilesCranmer

src/ChainRules.jl

@@ -42,43 +42,37 @@ function CRC.rrule(
         primal .= NaN
     end
 
-    # TODO: Preferable to use the primal in the pullback somehow
-    function pullback((dY, _))
-        dtree = let X = X, dY = dY, tree = tree, operators = operators
-            @thunk(
-                let
-                    _, gradient, complete2 = eval_grad_tree_array(
-                        tree, X, operators; variable=Val(false)
-                    )
-                    if !complete2
-                        gradient .= NaN
-                    end
+    return (primal, complete), EvalPullback(tree, X, operators)
+end
+
+# Wrap in struct rather than closure to ensure variables are boxed
+struct EvalPullback{N,A,O} <: Function
+    tree::N
+    X::A
+    operators::O
+end
 
-                    NodeTangent(
-                        tree,
-                        sum(j -> gradient[:, j] * dY[j], eachindex(dY, axes(gradient, 2))),
-                    )
-                end
-            )
-        end
-        dX = let X = X, dY = dY, tree = tree, operators = operators
-            @thunk(
-                let
-                    _, gradient2, complete3 = eval_grad_tree_array(
-                        tree, X, operators; variable=Val(true)
-                    )
-                    if !complete3
-                        gradient2 .= NaN
-                    end
+# TODO: Preferable to use the primal in the pullback somehow
+function (e::EvalPullback)((dY, _))
+    _, dX_constants_dY, complete = eval_grad_tree_array(
+        e.tree, e.X, e.operators; variable=Val(:both)
+    )
 
-                    gradient2 .* reshape(dY, 1, length(dY))
-                end
-            )
-        end
-        return (NoTangent(), dtree, dX, NoTangent())
+    if !complete
+        dX_constants_dY .= NaN
     end
 
-    return (primal, complete), pullback
+    nfeatures = size(e.X, 1)
+    dX_dY = @view dX_constants_dY[1:nfeatures, :]
+    dconstants_dY = @view dX_constants_dY[(nfeatures + 1):end, :]
+
+    dtree = NodeTangent(
+        e.tree, sum(j -> dconstants_dY[:, j] * dY[j], eachindex(dY, axes(dconstants_dY, 2)))
+    )
+
+    dX = dX_dY .* reshape(dY, 1, length(dY))
+
+    return (NoTangent(), dtree, dX, NoTangent())
 end
 
 end

src/EvaluateDerivative.jl

@@ -206,17 +206,29 @@ function eval_grad_tree_array(
     variable::Union{Bool,Val}=Val(false),
     turbo::Union{Bool,Val}=Val(false),
 ) where {T<:Number}
-    n_gradients = if isa(variable, Val{true}) || (isa(variable, Bool) && variable)
+    variable_mode = isa(variable, Val{true}) || (isa(variable, Bool) && variable)
+    constant_mode = isa(variable, Val{false}) || (isa(variable, Bool) && !variable)
+    both_mode = isa(variable, Val{:both})
+
+    n_gradients = if variable_mode
         size(cX, 1)::Int
-    else
+    elseif constant_mode
         count_constants(tree)::Int
+    elseif both_mode
+        size(cX, 1) + count_constants(tree)
     end
-    result = if isa(variable, Val{true}) || (variable isa Bool && variable)
+
+    result = if variable_mode
         eval_grad_tree_array(tree, n_gradients, nothing, cX, operators, Val(true))
-    else
+    elseif constant_mode
         index_tree = index_constants(tree)
         eval_grad_tree_array(tree, n_gradients, index_tree, cX, operators, Val(false))
+    elseif both_mode
+        # features come first because we can use size(cX, 1) to skip them
+        index_tree = index_constants(tree)
+        eval_grad_tree_array(tree, n_gradients, index_tree, cX, operators, Val(:both))
     end
+
     return (result.x, result.dx, result.ok)
 end
 
@@ -226,11 +238,9 @@ function eval_grad_tree_array(
     index_tree::Union{NodeIndex,Nothing},
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
-    ::Val{variable},
-)::ResultOk2 where {T<:Number,variable}
-    result = _eval_grad_tree_array(
-        tree, n_gradients, index_tree, cX, operators, Val(variable)
-    )
+    ::Val{mode},
+)::ResultOk2 where {T<:Number,mode}
+    result = _eval_grad_tree_array(tree, n_gradients, index_tree, cX, operators, Val(mode))
     !result.ok && return result
     return ResultOk2(
         result.x, result.dx, !(is_bad_array(result.x) || is_bad_array(result.dx))
@@ -260,30 +270,18 @@ end
     index_tree::Union{NodeIndex,Nothing},
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
-    ::Val{variable},
-)::ResultOk2 where {T<:Number,variable}
+    ::Val{mode},
+)::ResultOk2 where {T<:Number,mode}
     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),
+            tree, n_gradients, index_tree, cX, operators.unaops[i], operators, Val(mode)
         )
     end
     deg2_branch_skeleton = quote
         grad_deg2_eval(
-            tree,
-            n_gradients,
-            index_tree,
-            cX,
-            operators.binops[i],
-            operators,
-            Val(variable),
+            tree, n_gradients, index_tree, cX, operators.binops[i], operators, Val(mode)
         )
     end
     deg1_branch = if nuna > OPERATOR_LIMIT_BEFORE_SLOWDOWN
@@ -310,7 +308,7 @@ end
     end
     quote
         if tree.degree == 0
-            grad_deg0_eval(tree, n_gradients, index_tree, cX, Val(variable))
+            grad_deg0_eval(tree, n_gradients, index_tree, cX, Val(mode))
         elseif tree.degree == 1
             $deg1_branch
         else
@@ -324,8 +322,8 @@ function grad_deg0_eval(
     n_gradients,
     index_tree::Union{NodeIndex,Nothing},
     cX::AbstractMatrix{T},
-    ::Val{variable},
-)::ResultOk2 where {T<:Number,variable}
+    ::Val{mode},
+)::ResultOk2 where {T<:Number,mode}
     const_part = deg0_eval(tree, cX).x
 
     zero_mat = if isa(cX, Array)
@@ -334,15 +332,24 @@ function grad_deg0_eval(
         hcat([fill_similar(zero(T), cX, axes(cX, 2)) for _ in 1:n_gradients]...)'
     end
 
-    if variable == tree.constant
+    if (mode isa Bool && mode == tree.constant)
+        # No gradients at this leaf node
         return ResultOk2(const_part, zero_mat, true)
     end
 
-    index = if variable
-        tree.feature
-    else
+    index = if (mode isa Bool && mode)
+        tree.feature::UInt16
+    elseif (mode isa Bool && !mode)
         (index_tree === nothing ? zero(UInt16) : index_tree.val::UInt16)
+    elseif mode == :both
+        index_tree::NodeIndex
+        if tree.constant
+            index_tree.val::UInt16 + UInt16(size(cX, 1))
+        else
+            tree.feature::UInt16
+        end
     end
+
     derivative_part = zero_mat
     derivative_part[index, :] .= one(T)
     return ResultOk2(const_part, derivative_part, true)
@@ -355,15 +362,15 @@ function grad_deg1_eval(
     cX::AbstractMatrix{T},
     op::F,
     operators::OperatorEnum,
-    ::Val{variable},
-)::ResultOk2 where {T<:Number,F,variable}
+    ::Val{mode},
+)::ResultOk2 where {T<:Number,F,mode}
     result = eval_grad_tree_array(
         tree.l,
         n_gradients,
         index_tree === nothing ? index_tree : index_tree.l,
         cX,
         operators,
-        Val(variable),
+        Val(mode),
     )
     !result.ok && return result
 
@@ -389,15 +396,15 @@ function grad_deg2_eval(
     cX::AbstractMatrix{T},
     op::F,
     operators::OperatorEnum,
-    ::Val{variable},
-)::ResultOk2 where {T<:Number,F,variable}
+    ::Val{mode},
+)::ResultOk2 where {T<:Number,F,mode}
     result_l = eval_grad_tree_array(
         tree.l,
         n_gradients,
         index_tree === nothing ? index_tree : index_tree.l,
         cX,
         operators,
-        Val(variable),
+        Val(mode),
     )
     !result_l.ok && return result_l
     result_r = eval_grad_tree_array(
@@ -406,7 +413,7 @@ function grad_deg2_eval(
         index_tree === nothing ? index_tree : index_tree.r,
         cX,
         operators,
-        Val(variable),
+        Val(mode),
     )
     !result_r.ok && return result_r