Merge pull request #88 from SymbolicML/fix-parametric-gradients

Fix gradients of parametric expressions

Author: MilesCranmer

src/ChainRules.jl

@@ -12,7 +12,7 @@ struct NodeTangent{T,N<:AbstractExpressionNode{T},A<:AbstractArray{T}} <: Abstra
     gradient::A
 end
 function Base.:+(a::NodeTangent, b::NodeTangent)
-    @assert a.tree == b.tree
+    # @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)

src/ParametricExpression.jl

@@ -1,10 +1,12 @@
 module ParametricExpressionModule
 
 using DispatchDoctor: @stable, @unstable
+using ChainRulesCore: ChainRulesCore, NoTangent, @thunk
 
-using ..OperatorEnumModule: AbstractOperatorEnum
+using ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum
 using ..NodeModule: AbstractExpressionNode, Node, tree_mapreduce
 using ..ExpressionModule: AbstractExpression, Metadata
+using ..ChainRulesModule: NodeTangent
 
 import ..NodeModule: constructorof, preserve_sharing, leaf_copy, leaf_hash, leaf_equal
 import ..NodeUtilsModule:
@@ -236,6 +238,36 @@ function Base.convert(::Type{Node}, ex::ParametricExpression{T}) where {T}
         Node{T},
     )
 end
+function ChainRulesCore.rrule(
+    ::typeof(convert), ::Type{Node}, ex::ParametricExpression{T}
+) where {T}
+    tree = get_contents(ex)
+    primal = convert(Node, ex)
+    pullback = let tree = tree
+        d_primal -> let
+            # ^The exact same tangent with respect to constants, so we can just take it.
+            d_ex = @thunk(
+                let
+                    parametric_node_tangent = NodeTangent(tree, d_primal.gradient)
+                    (;
+                        tree=parametric_node_tangent,
+                        metadata=(;
+                            _data=(;
+                                operators=NoTangent(),
+                                variable_names=NoTangent(),
+                                parameters=NoTangent(),
+                                parameter_names=NoTangent(),
+                            )
+                        ),
+                    )
+                end
+            )
+            (NoTangent(), NoTangent(), d_ex)
+        end
+    end
+    return primal, pullback
+end
+
 #! format: off
 function (ex::ParametricExpression)(X::AbstractMatrix, operators::Union{AbstractOperatorEnum,Nothing}=nothing; kws...)
     return eval_tree_array(ex, X, operators; kws...)  # Will error
@@ -250,7 +282,7 @@ function (ex::ParametricExpression)(
     operators::Union{AbstractOperatorEnum,Nothing}=nothing;
     kws...,
 ) where {T}
-    (output, flag) = eval_tree_array(ex, X, classes, operators; kws...)  # Will error
+    (output, flag) = eval_tree_array(ex, X, classes, operators; kws...)
     if !flag
         output .= NaN
     end
@@ -276,6 +308,7 @@ function eval_tree_array(
     regular_tree = convert(Node, ex)
     return eval_tree_array(regular_tree, params_and_X, get_operators(ex, operators); kws...)
 end
+
 function string_tree(
     ex::ParametricExpression,
     operators::Union{AbstractOperatorEnum,Nothing}=nothing;

test/Project.toml

@@ -3,6 +3,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

test/test_parametric_expression.jl

@@ -266,3 +266,54 @@ end
         )
     end
 end
+
+@testitem "Parametric expression derivatives" begin
+    using DynamicExpressions
+    using Zygote: Zygote
+    using Random: MersenneTwister
+    using DifferentiationInterface: value_and_gradient, AutoZygote
+
+    rng = MersenneTwister(0)
+    X = rand(rng, 2, 32)
+    true_params = [0.5 2.0]
+    init_params = [0.1 0.2]
+    classes = rand(rng, 1:2, 32)
+    y = [X[1, i] * X[1, i] - cos(2.6 * X[2, i]) + true_params[1, classes[i]] for i in 1:32]
+
+    (true_val, true_grad) =
+        value_and_gradient(AutoZygote(), (X, init_params, [2.5])) do (X, params, c)
+            pred = [
+                X[1, i] * X[1, i] - cos(c[1] * X[2, i]) + params[1, classes[i]] for
+                i in 1:32
+            ]
+            sum(abs2, pred .- y)
+        end
+
+    operators = OperatorEnum(; unary_operators=[cos], binary_operators=[+, *, -])
+    ex = @parse_expression(
+        x * x - cos(2.5 * y) + p1,
+        operators = operators,
+        expression_type = ParametricExpression,
+        variable_names = ["x", "y"],
+        extra_metadata = (parameter_names=["p1"], parameters=init_params)
+    )
+    f = let operators = operators, X = X, classes = classes, y = y
+        ex -> sum(abs2, ex(X, classes) .- y)
+    end
+    @test f(ex) isa Float64
+    (val, grad) = value_and_gradient(f, AutoZygote(), ex)
+
+    @test val isa Float64
+    @test grad isa NamedTuple
+    @test grad.tree isa DynamicExpressions.ChainRulesModule.NodeTangent{
+        Float64,ParametricNode{Float64},Vector{Float64}
+    }
+    @test grad.metadata._data.parameters isa Matrix{Float64}
+
+    # Loss value:
+    @test val ≈ true_val
+    # Gradient w.r.t. the constant:
+    @test grad.tree.gradient ≈ true_grad[3]
+    # Gradient w.r.t. the parameters:
+    @test grad.metadata._data.parameters ≈ true_grad[2]
+end