feat: make diff compatibility with n-arity

Author: MilesCranmer

ext/DynamicExpressionsZygoteExt.jl

@@ -3,30 +3,17 @@ module DynamicExpressionsZygoteExt
 using Zygote: gradient
 import DynamicExpressions.ExtensionInterfaceModule: _zygote_gradient, ZygoteGradient
 
-function _zygote_gradient(op::F, ::Val{1}) where {F}
-    return ZygoteGradient{F,1,1}(op)
-end
-function _zygote_gradient(op::F, ::Val{2}, ::Val{side}=Val(nothing)) where {F,side}
-    # side should be either nothing (for both), 1, or 2
-    @assert side === nothing || side in (1, 2)
-    return ZygoteGradient{F,2,side}(op)
+function _zygote_gradient(op::F, ::Val{degree}) where {F,degree}
+    return ZygoteGradient{F,degree}(op)
 end
 
-function (g::ZygoteGradient{F,1,1})(x) where {F}
+function (g::ZygoteGradient{F,1})(x) where {F}
     out = only(gradient(g.op, x))
     return out === nothing ? zero(x) : out
 end
-function (g::ZygoteGradient{F,2,nothing})(x, y) where {F}
-    (∂x, ∂y) = gradient(g.op, x, y)
-    return (∂x === nothing ? zero(x) : ∂x, ∂y === nothing ? zero(y) : ∂y)
-end
-function (g::ZygoteGradient{F,2,1})(x, y) where {F}
-    ∂x = only(gradient(Base.Fix2(g.op, y), x))
-    return ∂x === nothing ? zero(x) : ∂x
-end
-function (g::ZygoteGradient{F,2,2})(x, y) where {F}
-    ∂y = only(gradient(Base.Fix1(g.op, x), y))
-    return ∂y === nothing ? zero(y) : ∂y
+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))
 end
 
 end

src/EvaluateDerivative.jl

@@ -1,6 +1,6 @@
 module EvaluateDerivativeModule
 
-import ..NodeModule: AbstractExpressionNode, constructorof
+import ..NodeModule: AbstractExpressionNode, constructorof, get_children
 import ..OperatorEnumModule: OperatorEnum
 import ..UtilsModule: fill_similar, ResultOk2
 import ..ValueInterfaceModule: is_valid_array
@@ -66,54 +66,18 @@ function eval_diff_tree_array(
 end
 
 @generated function _eval_diff_tree_array(
-    tree::AbstractExpressionNode{T},
+    tree::AbstractExpressionNode{T,D},
     cX::AbstractMatrix{T},
     operators::OperatorEnum,
     direction::Integer,
-)::ResultOk2 where {T<:Number}
-    nuna = get_nuna(operators)
-    nbin = get_nbin(operators)
-    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)
-            )
-        end
-    end
-    deg2_branch = if nbin > OPERATOR_LIMIT_BEFORE_SLOWDOWN
-        quote
-            diff_deg2_eval(tree, cX, operators.binops[op_idx], operators, direction)
-        end
-    else
-        quote
-            Base.Cartesian.@nif(
-                $nbin,
-                i -> i == op_idx,
-                i ->
-                    diff_deg2_eval(tree, cX, operators.binops[i], operators, direction)
-            )
-        end
-    end
+)::ResultOk2 where {T<:Number,D}
     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_valid_array(result.x) && is_valid_array(result.dx)
+        deg = tree.degree
+        deg == 0 && return diff_deg0_eval(tree, cX, direction)
+        Base.Cartesian.@nif(
+            $D,
+            i -> i == deg,
+            i -> dispatch_diff_degn_eval(tree, cX, Val(i), operators, direction)
         )
     end
 end
@@ -130,58 +94,71 @@ function diff_deg0_eval(
     return ResultOk2(const_part, derivative_part, true)
 end
 
-function diff_deg1_eval(
-    tree::AbstractExpressionNode{T},
-    cX::AbstractMatrix{T},
-    op::F,
-    operators::OperatorEnum,
-    direction::Integer,
-) where {T<:Number,F}
-    result = _eval_diff_tree_array(tree.l, cX, operators, direction)
-    !result.ok && return result
-
-    # TODO - add type assertions to get better speed:
-    cumulator = result.x
-    dcumulator = result.dx
-    diff_op = _zygote_gradient(op, Val(1))
-    @inbounds @simd for j in eachindex(cumulator)
-        x = op(cumulator[j])::T
-        dx = diff_op(cumulator[j])::T * dcumulator[j]
-
-        cumulator[j] = x
-        dcumulator[j] = dx
+@generated function diff_degn_eval(
+    x_cumulators::NTuple{N}, dx_cumulators::NTuple{N}, op::F, direction::Integer
+) where {N,F}
+    quote
+        Base.Cartesian.@nexprs($N, i -> begin
+            x_cumulator_i = x_cumulators[i]
+            dx_cumulator_i = dx_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]
+            )
+            dx = Base.Cartesian.@ncall($N, +, i -> grad_i * dx_cumulator_i[j])
+            x_cumulator_1[j] = x
+            dx_cumulator_1[j] = dx
+        end
+        return ResultOk2(x_cumulator_1, dx_cumulator_1, true)
     end
-    return result
 end
 
-function diff_deg2_eval(
-    tree::AbstractExpressionNode{T},
+@generated function dispatch_diff_degn_eval(
+    tree::AbstractExpressionNode{T,D},
     cX::AbstractMatrix{T},
-    op::F,
-    operators::OperatorEnum,
+    ::Val{degree},
+    operators::OperatorEnum{OPS},
     direction::Integer,
-) where {T<:Number,F}
-    result_l = _eval_diff_tree_array(tree.l, cX, operators, direction)
-    !result_l.ok && return result_l
-    result_r = _eval_diff_tree_array(tree.r, cX, operators, direction)
-    !result_r.ok && return result_r
-
-    ar_l = result_l.x
-    d_ar_l = result_l.dx
-    ar_r = result_r.x
-    d_ar_r = result_r.dx
-    diff_op = _zygote_gradient(op, Val(2))
-
-    @inbounds @simd for j in eachindex(ar_l)
-        x = op(ar_l[j], ar_r[j])::T
-
-        first, second = diff_op(ar_l[j], ar_r[j])::Tuple{T,T}
-        dx = first * d_ar_l[j] + second * d_ar_r[j]
+) where {T<:Number,D,degree,OPS}
+    nops = length(OPS.types[degree].types)
+
+    setup = quote
+        cs = get_children(tree, Val($degree))
+        Base.Cartesian.@nexprs(
+            $degree,
+            i -> begin
+                result_i = _eval_diff_tree_array(cs[i], cX, operators, direction)
+                !result_i.ok && return result_i
+            end
+        )
+        x_cumulators = Base.Cartesian.@ntuple($degree, i -> result_i.x)
+        dx_cumulators = Base.Cartesian.@ntuple($degree, i -> result_i.dx)
+        op_idx = tree.op
+    end
 
-        ar_l[j] = x
-        d_ar_l[j] = dx
+    if nops > OPERATOR_LIMIT_BEFORE_SLOWDOWN
+        quote
+            $setup
+            diff_degn_eval(
+                x_cumulators, dx_cumulators, operators[$degree][op_idx], direction
+            )
+        end
+    else
+        quote
+            $setup
+            Base.Cartesian.@nif(
+                $nops,
+                i -> i == op_idx,
+                i -> diff_degn_eval(
+                    x_cumulators, dx_cumulators, operators[$degree][i], direction
+                )
+            )
+        end
     end
-    return result_l
+    # TODO: Need to add the case for many operators
 end
 
 """

src/ExtensionInterface.jl

@@ -7,19 +7,12 @@ function symbolic_to_node(args...; kws...)
     return error("Please load the `SymbolicUtils` package to use `symbolic_to_node`.")
 end
 
-struct ZygoteGradient{F,degree,arg} <: Function
+struct ZygoteGradient{F,degree} <: Function
     op::F
 end
 
-function Base.show(io::IO, g::ZygoteGradient{F,degree,arg}) where {F,degree,arg}
+function Base.show(io::IO, g::ZygoteGradient{F,degree}) where {F,degree}
     print(io, "∂")
-    if degree == 2
-        if arg == 1
-            print(io, "₁")
-        elseif arg == 2
-            print(io, "₂")
-        end
-    end
     print(io, g.op)
     return nothing
 end

src/Node.jl

@@ -6,15 +6,15 @@ import ..OperatorEnumModule: AbstractOperatorEnum
 import ..UtilsModule: deprecate_varmap, Undefined
 
 const DEFAULT_NODE_TYPE = Float32
+const DEFAULT_MAX_DEGREE = 2
 
 """
     AbstractNode{D}
 
 Abstract type for D-arity trees. Must have the following fields:
 
-- `degree::Integer`: Degree of the node. Either 0, 1, or 2. If 1,
-    then `l` needs to be defined as the left child. If 2,
-    then `r` also needs to be defined as the right child.
+- `degree::UInt8`: Degree of the node. This should be a value
+    between 0 and `DEFAULT_MAX_DEGREE`.
 - `children`: A collection of D references to children nodes.
 
 # Deprecated fields
@@ -25,7 +25,7 @@ Abstract type for D-arity trees. Must have the following fields:
     Don't use `nothing` to represent an undefined value
     as it will incur a large performance penalty.
 - `r::AbstractNode{D}`: Right child of the current node. Should only
-    be defined if `degree == 2`.
+    be defined if `degree >= 2`.
 """
 abstract type AbstractNode{D} end
 
@@ -82,7 +82,7 @@ for N in (:Node, :GraphNode)
         ## Constructors:
         #################
         $N{_T,_D}() where {_T,_D} = new{_T,_D::Int}()
-        $N{_T}() where {_T} = $N{_T,2}()
+        $N{_T}() where {_T} = $N{_T,DEFAULT_MAX_DEGREE}()
         # TODO: Test with this disabled to spot any unintended uses
     end
 end
@@ -250,7 +250,6 @@ end
 Base.eltype(::Type{<:AbstractExpressionNode{T}}) where {T} = T
 Base.eltype(::AbstractExpressionNode{T}) where {T} = T
 
-const DEFAULT_MAX_DEGREE = 2
 max_degree(::Type{<:AbstractNode}) = DEFAULT_MAX_DEGREE
 max_degree(::Type{<:AbstractNode{D}}) where {D} = D
 max_degree(node::AbstractNode) = max_degree(typeof(node))

test/test_zygote_gradient_wrapper.jl

@@ -11,23 +11,8 @@
     g(x, y) = x * y
     @test repr(_zygote_gradient(g, Val(2))) == "∂g"
 
-    # Test binary gradient (first partial)
-    @test repr(_zygote_gradient(g, Val(2), Val(1))) == "∂₁g"
-
-    # Test binary gradient (second partial)
-    @test repr(_zygote_gradient(g, Val(2), Val(2))) == "∂₂g"
-
     # Test with standard operators
     @test repr(_zygote_gradient(+, Val(2))) == "∂+"
-    @test repr(_zygote_gradient(*, Val(2), Val(1))) == "∂₁*"
-    @test repr(_zygote_gradient(*, Val(2), Val(2))) == "∂₂*"
-
-    first_partial = _zygote_gradient(log, Val(2), Val(1))
-    nested = _zygote_gradient(first_partial, Val(1))
-    @test repr(nested) == "∂∂₁log"
-
-    # Also should work with text/plain
-    @test repr("text/plain", nested) == "∂∂₁log"
 end
 
 @testitem "ZygoteGradient evaluation" begin
@@ -45,10 +30,4 @@ end
     # Test binary gradient (both partials)
     g(x, y) = x * y
     @test (_zygote_gradient(g, Val(2)))(x, y) == (3.0, 2.0)
-
-    # Test binary gradient (first partial)
-    @test (_zygote_gradient(g, Val(2), Val(1)))(x, y) == 3.0
-
-    # Test second partial
-    @test (_zygote_gradient(g, Val(2), Val(2)))(x, y) == 2.0
 end