test: gradients of parametric expression

Author: MilesCranmer

src/ParametricExpression.jl

@@ -293,6 +293,7 @@ function ChainRulesCore.rrule(
         parameters = ex.metadata.parameters
         num_params = size(parameters, 1)
         num_classes = size(parameters, 2)
+        _operators = get_operators(ex, operators)
         indexed_parameters = [
             parameters[i_parameter, classes[i_row]] for
             i_parameter in eachindex(axes(parameters, 1)), i_row in eachindex(classes)
@@ -302,10 +303,10 @@ function ChainRulesCore.rrule(
         regular_tree = convert(Node, ex)
 
         _, gradient_tree, complete1 = eval_grad_tree_array(
-            regular_tree, params_and_X, operators; variable=Val(false)
+            regular_tree, params_and_X, _operators; variable=Val(false)
         )
         _, gradient_params_and_X, complete2 = eval_grad_tree_array(
-            regular_tree, params_and_X, operators; variable=Val(true)
+            regular_tree, params_and_X, _operators; variable=Val(true)
         )
 
         if !complete1
@@ -331,13 +332,15 @@ function ChainRulesCore.rrule(
         d_ex = (;
             tree=d_tree,
             metadata=(;
-                operators=NoTangent(),
-                variable_names=NoTangent(),
-                parameters=d_parameters,
-                parameter_names=NoTangent(),
+                _data=(;
+                    operators=NoTangent(),
+                    variable_names=NoTangent(),
+                    parameters=d_parameters,
+                    parameter_names=NoTangent(),
+                ),
             ),
         )
-        return (NoTangent(), d_ex, copy(d_X), NoTangent(), NoTangent())
+        return (NoTangent(), d_ex, d_X, NoTangent(), NoTangent())
     end
 
     return (primal, complete), pullback

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,37 @@ 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)
+    classes = rand(rng, 1:2, 32)
+    y = @. X[1, :] * X[1, :] - cos(2.6 * X[2, :]) + classes
+
+    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=[0.5 0.2])
+    )
+    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}
+end