fix: extract_gradient for regular expressions

Author: MilesCranmer

src/Expression.jl

@@ -272,13 +272,7 @@ function set_constants!(ex::Expression{T}, constants, refs) where {T}
     return set_constants!(get_tree(ex), constants, refs)
 end
 function extract_gradient(
-    gradient::@NamedTuple{
-        tree::NT,
-        metadata::@NamedTuple{
-            _data::@NamedTuple{operators::Nothing, variable_names::Nothing}
-        }
-    },
-    ex::Expression{T,N},
+    gradient::@NamedTuple{tree::NT, metadata::Nothing}, ex::Expression{T,N}
 ) where {T,N<:AbstractExpressionNode{T},NT<:NodeTangent{T,N}}
     # TODO: This messy gradient type is produced by ChainRules. There is probably a better way to do this.
     return extract_gradient(gradient.tree, get_tree(ex))

test/test_expressions.jl

@@ -76,6 +76,19 @@ end
     end
 end
 
+@testitem "Can also get derivatives of expression itself" begin
+    using DynamicExpressions
+    using Zygote: Zygote
+    using DifferentiationInterface: AutoZygote, gradient
+
+    ex = @parse_expression(x1 + 1.5, binary_operators = [+], variable_names = ["x1"])
+    d_ex = gradient(AutoZygote(), ex) do ex
+        sum(ex(ones(1, 5)))
+    end
+    @test d_ex isa NamedTuple
+    @test extract_gradient(d_ex, ex) ≈ [5.0]
+end
+
 @testitem "Expression simplification" begin
     using DynamicExpressions