feat: make grad compatible with n-arity

Author: MilesCranmer

ext/DynamicExpressionsZygoteExt.jl

@@ -7,10 +7,7 @@ function _zygote_gradient(op::F, ::Val{degree}) where {F,degree}
     return ZygoteGradient{F,degree}(op)
 end
 
-function (g::ZygoteGradient{F,1})(x) where {F}
-    out = only(gradient(g.op, x))
-    return out === nothing ? zero(x) : out
-end
+# All this does is remove `nothing`, so that we get type stability
 function (g::ZygoteGradient{F,degree})(args::Vararg{Any,degree}) where {F,degree}
     partials = gradient(g.op, args...)
     return ntuple(i -> @something(partials[i], zero(args[i])), Val(degree))

src/EvaluateDerivative.jl

@@ -254,55 +254,97 @@ function eval_grad_tree_array(
 end
 
 @generated function _eval_grad_tree_array(
-    tree::AbstractExpressionNode{T},
+    tree::AbstractExpressionNode{T,D},
     n_gradients,
-    index_tree::Union{NodeIndex,Nothing},
+    index_tree::Union{NodeIndex{<:Any,D},Nothing},
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
     ::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(mode)
+)::ResultOk2 where {T<:Number,D,mode}
+    quote
+        deg = tree.degree
+        deg == 0 && return grad_deg0_eval(tree, n_gradients, index_tree, cX, Val(mode))
+        Base.Cartesian.@nif(
+            $D,
+            i -> i == deg,
+            i -> dispatch_grad_degn_eval(
+                tree, n_gradients, index_tree, cX, Val(i), operators, Val(mode)
+            )
         )
     end
-    deg2_branch_skeleton = quote
-        grad_deg2_eval(
-            tree, n_gradients, index_tree, cX, operators.binops[i], operators, Val(mode)
+end
+
+@generated function dispatch_grad_degn_eval(
+    tree::AbstractExpressionNode{T},
+    n_gradients,
+    index_tree::Union{NodeIndex,Nothing},
+    cX::AbstractMatrix{T},
+    ::Val{degree},
+    operators::OperatorEnum{OPS},
+    ::Val{mode},
+) where {T<:Number,degree,OPS,mode}
+    setup = quote
+        cs = get_children(tree, Val($degree))
+        index_cs =
+            isnothing(index_tree) ? index_tree : get_children(index_tree, Val($degree))
+        Base.Cartesian.@nexprs(
+            $degree,
+            i -> begin
+                result_i = eval_grad_tree_array(
+                    cs[i],
+                    n_gradients,
+                    isnothing(index_cs) ? index_cs : index_cs[i],
+                    cX,
+                    operators,
+                    Val(mode),
+                )
+                !result_i.ok && return result_i
+            end
         )
+        x_cumulators = Base.Cartesian.@ntuple($degree, i -> result_i.x)
+        d_cumulators = Base.Cartesian.@ntuple($degree, i -> result_i.dx)
+        op_idx = tree.op
     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
+    nops = length(OPS.types[degree].types)
+    if nops > OPERATOR_LIMIT_BEFORE_SLOWDOWN
         quote
-            i = tree.op
-            $deg2_branch_skeleton
+            $setup
+            grad_degn_eval(x_cumulators, d_cumulators, operators[$degree][op_idx])
         end
     else
         quote
-            op_idx = tree.op
-            Base.Cartesian.@nif($nbin, i -> i == op_idx, i -> $deg2_branch_skeleton)
+            $setup
+            Base.Cartesian.@nif(
+                $nops,
+                i -> i == op_idx,
+                i -> grad_degn_eval(x_cumulators, d_cumulators, operators[$degree][i])
+            )
         end
     end
+end
+
+@generated function grad_degn_eval(
+    x_cumulators::NTuple{N}, d_cumulators::NTuple{N}, op::F
+) where {N,F}
     quote
-        if tree.degree == 0
-            grad_deg0_eval(tree, n_gradients, index_tree, cX, Val(mode))
-        elseif tree.degree == 1
-            $deg1_branch
-        else
-            $deg2_branch
+        Base.Cartesian.@nexprs($N, i -> begin
+            x_cumulator_i = x_cumulators[i]
+            d_cumulator_i = d_cumulators[i]
+        end)
+        diff_op = _zygote_gradient(op, Val($N))
+        @inbounds @simd for j in eachindex(x_cumulator_1)
+            x = Base.Cartesian.@ncall($N, op, i -> x_cumulator_i[j])
+            Base.Cartesian.@ntuple($N, i -> grad_i) = Base.Cartesian.@ncall(
+                $N, diff_op, i -> x_cumulator_i[j]
+            )
+            x_cumulator_1[j] = x
+            for k in axes(d_cumulator_1, 1)
+                d_cumulator_1[k, j] = Base.Cartesian.@ncall(
+                    $N, +, i -> grad_i * d_cumulator_i[k, j]
+                )
+            end
         end
+        return ResultOk2(x_cumulator_1, d_cumulator_1, true)
     end
 end
 
@@ -344,85 +386,4 @@ function grad_deg0_eval(
     return ResultOk2(const_part, derivative_part, true)
 end
 
-function grad_deg1_eval(
-    tree::AbstractExpressionNode{T},
-    n_gradients,
-    index_tree::Union{NodeIndex,Nothing},
-    cX::AbstractMatrix{T},
-    op::F,
-    operators::OperatorEnum,
-    ::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(mode),
-    )
-    !result.ok && return result
-
-    cumulator = result.x
-    dcumulator = result.dx
-    diff_op = _zygote_gradient(op, Val(1))
-    @inbounds @simd for j in axes(dcumulator, 2)
-        x = op(cumulator[j])::T
-        dx = diff_op(cumulator[j])::T
-
-        cumulator[j] = x
-        for k in axes(dcumulator, 1)
-            dcumulator[k, j] = dx * dcumulator[k, j]
-        end
-    end
-    return result
-end
-
-function grad_deg2_eval(
-    tree::AbstractExpressionNode{T},
-    n_gradients,
-    index_tree::Union{NodeIndex,Nothing},
-    cX::AbstractMatrix{T},
-    op::F,
-    operators::OperatorEnum,
-    ::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(mode),
-    )
-    !result_l.ok && return result_l
-    result_r = eval_grad_tree_array(
-        tree.r,
-        n_gradients,
-        index_tree === nothing ? index_tree : index_tree.r,
-        cX,
-        operators,
-        Val(mode),
-    )
-    !result_r.ok && return result_r
-
-    cumulator_l = result_l.x
-    dcumulator_l = result_l.dx
-    cumulator_r = result_r.x
-    dcumulator_r = result_r.dx
-    diff_op = _zygote_gradient(op, Val(2))
-    @inbounds @simd for j in axes(dcumulator_l, 2)
-        c1 = cumulator_l[j]
-        c2 = cumulator_r[j]
-        x = op(c1, c2)::T
-        dx1, dx2 = diff_op(c1, c2)::Tuple{T,T}
-        cumulator_l[j] = x
-        for k in axes(dcumulator_l, 1)
-            dcumulator_l[k, j] = dx1 * dcumulator_l[k, j] + dx2 * dcumulator_r[k, j]
-        end
-    end
-
-    return result_l
-end
-
 end

test/test_zygote_gradient_wrapper.jl

@@ -25,7 +25,7 @@ end
 
     # Test unary gradient
     f(x) = x^2
-    @test (_zygote_gradient(f, Val(1)))(x) == 4.0
+    @test (_zygote_gradient(f, Val(1)))(x) == (4.0,)
 
     # Test binary gradient (both partials)
     g(x, y) = x * y