feat: add ChainRulesCore.rrule for eval_tree_array

Author: MilesCranmer

Project.toml

@@ -4,6 +4,7 @@ authors = ["MilesCranmer <miles.cranmer@gmail.com>"]
 version = "0.16.0"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"

src/ChainRules.jl

@@ -0,0 +1,75 @@
+module ChainRulesModule
+
+using ChainRulesCore:
+    ChainRulesCore, AbstractTangent, NoTangent, ZeroTangent, Tangent, @thunk, canonicalize
+using ..OperatorEnumModule: OperatorEnum
+using ..EquationModule: AbstractExpressionNode, with_type_parameters, tree_mapreduce
+using ..EvaluateEquationModule: eval_tree_array
+using ..EvaluateEquationDerivativeModule: eval_grad_tree_array
+
+struct NodeTangent{T,N<:AbstractExpressionNode{T},A<:AbstractArray{T}} <: AbstractTangent
+    tree::N
+    gradient::A
+end
+function Base.:+(a::NodeTangent, b::NodeTangent)
+    @assert a.tree == b.tree
+    return NodeTangent(a.tree, a.gradient + b.gradient)
+end
+Base.:*(a::Number, b::NodeTangent) = NodeTangent(b.tree, a * b.gradient)
+Base.:*(a::NodeTangent, b::Number) = NodeTangent(a.tree, a.gradient * b)
+Base.zero(::Union{Type{NodeTangent},NodeTangent}) = ZeroTangent()
+
+function ChainRulesCore.rrule(
+    ::typeof(eval_tree_array),
+    tree::AbstractExpressionNode,
+    X::AbstractMatrix,
+    operators::OperatorEnum;
+    turbo=Val(false),
+    bumper=Val(false),
+)
+    primal, complete = eval_tree_array(tree, X, operators; turbo, bumper)
+
+    if !complete
+        primal .= NaN
+    end
+
+    # TODO: Preferable to use the primal in the pullback somehow
+    function pullback((dY, _))
+        dtree = let dY = dY, tree = tree, operators = operators
+            @thunk(
+                let
+                    _, gradient, complete = eval_grad_tree_array(
+                        tree, X, operators; variable=Val(false)
+                    )
+                    if !complete
+                        gradient .= NaN
+                    end
+
+                    NodeTangent(
+                        tree,
+                        sum(j -> gradient[:, j] * dY[j], eachindex(dY, axes(gradient, 2))),
+                    )
+                end
+            )
+        end
+        dX = let dY = dY, tree = tree, operators = operators
+            @thunk(
+                let
+                    _, gradient, complete = eval_grad_tree_array(
+                        tree, X, operators; variable=Val(true)
+                    )
+                    if !complete
+                        gradient .= NaN
+                    end
+
+                    gradient .* reshape(dY, 1, length(dY))
+                end
+            )
+        end
+        return (NoTangent(), dtree, dX, NoTangent())
+    end
+
+    return (primal, complete), pullback
+end
+
+end

src/DynamicExpressions.jl

@@ -8,6 +8,7 @@ include("EquationUtils.jl")
 include("Strings.jl")
 include("EvaluateEquation.jl")
 include("EvaluateEquationDerivative.jl")
+include("ChainRules.jl")
 include("EvaluationHelpers.jl")
 include("SimplifyEquation.jl")
 include("OperatorEnumConstruction.jl")