feat: introduce generic extract_gradients for flattening gradients

Author: MilesCranmer

src/ChainRules.jl

@@ -17,6 +17,9 @@ struct NodeTangent{T,N<:AbstractExpressionNode{T},A<:AbstractArray{T}} <: Abstra
     tree::N
     gradient::A
 end
+function extract_gradient(gradient::NodeTangent, ::AbstractExpressionNode)
+    return gradient.gradient
+end
 function Base.:+(a::NodeTangent, b::NodeTangent)
     # @assert a.tree == b.tree
     return NodeTangent(a.tree, a.gradient + b.gradient)

src/DynamicExpressions.jl

@@ -61,7 +61,7 @@ import .NodeModule:
     OperatorEnum, GenericOperatorEnum, @extend_operators, set_default_variable_names!
 @reexport import .EvaluateModule: eval_tree_array, differentiable_eval_tree_array
 @reexport import .EvaluateDerivativeModule: eval_diff_tree_array, eval_grad_tree_array
-@reexport import .ChainRulesModule: NodeTangent
+@reexport import .ChainRulesModule: NodeTangent, extract_gradient
 @reexport import .SimplifyModule: combine_operators, simplify_tree!
 @reexport import .EvaluationHelpersModule
 @reexport import .ExtensionInterfaceModule: node_to_symbolic, symbolic_to_node

src/Expression.jl

@@ -5,6 +5,7 @@ using DispatchDoctor: @unstable
 using ..NodeModule: AbstractExpressionNode, Node
 using ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum
 using ..UtilsModule: Undefined
+using ..ChainRulesModule: NodeTangent
 
 import ..NodeModule: copy_node, set_node!, count_nodes, tree_mapreduce, constructorof
 import ..NodeUtilsModule:
@@ -16,6 +17,7 @@ import ..NodeUtilsModule:
     has_constants,
     get_constants,
     set_constants!
+import ..ChainRulesModule: extract_gradient
 
 """A wrapper for a named tuple to avoid piracy."""
 struct Metadata{NT<:NamedTuple}
@@ -140,6 +142,12 @@ end
 function set_constants!(ex::AbstractExpression{T}, constants, refs) where {T}
     return error("`set_constants!` function must be implemented for $(typeof(ex)) types.")
 end
+function extract_gradient(gradient, ex::AbstractExpression)
+    # Should match `get_constants`
+    return error(
+        "`extract_gradient` function must be implemented for $(typeof(ex)) types with $(typeof(gradient)) gradient.",
+    )
+end
 function get_contents(ex::AbstractExpression)
     return error("`get_contents` function must be implemented for $(typeof(ex)) types.")
 end
@@ -263,6 +271,18 @@ end
 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},
+) 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))
+end
 
 import ..StringsModule: string_tree, print_tree
 

src/Interfaces.jl

@@ -159,6 +159,7 @@ ei_components = (
         default_node_type = "returns the default node type for the expression" => _check_default_node,
         constructorof = "gets the constructor function for a type" => _check_constructorof,
         tree_mapreduce = "applies a function across the tree" => _check_tree_mapreduce
+        # TODO: add extract_gradient(gradient, ex::AbstractExpression)
     )
 )
 ei_description = (

src/ParametricExpression.jl

@@ -20,6 +20,7 @@ import ..StringsModule: string_tree
 import ..EvaluateModule: eval_tree_array
 import ..EvaluateDerivativeModule: eval_grad_tree_array
 import ..EvaluationHelpersModule: _grad_evaluator
+import ..ChainRulesModule: extract_gradient
 import ..ExpressionModule:
     get_contents,
     get_metadata,
@@ -207,7 +208,7 @@ has_constants(ex::ParametricExpression) = _interface_error()
 has_operators(ex::ParametricExpression) = has_operators(get_tree(ex))
 function get_constants(ex::ParametricExpression{T}) where {T}
     constants, constant_refs = get_constants(get_tree(ex))
-    parameters = ex.metadata.parameters
+    parameters = get_metadata(ex).parameters
     flat_parameters = parameters[:]
     num_constants = length(constants)
     num_parameters = length(flat_parameters)
@@ -218,9 +219,27 @@ function set_constants!(ex::ParametricExpression{T}, x, refs) where {T}
     # First, set the usual constants
     set_constants!(get_tree(ex), @view(x[1:(refs.num_constants)]), refs.constant_refs)
     # Then, copy in the parameters
-    ex.metadata.parameters[:] .= @view(x[(refs.num_constants + 1):end])
+    get_metadata(ex).parameters[:] .= @view(x[(refs.num_constants + 1):end])
     return ex
 end
+function extract_gradient(
+    gradient::@NamedTuple{
+        tree::NT,
+        metadata::@NamedTuple{
+            _data::@NamedTuple{
+                operators::Nothing,
+                variable_names::Nothing,
+                parameters::PARAM,
+                parameter_names::Nothing,
+            }
+        }
+    },
+    ex::ParametricExpression{T,N},
+) where {T,N<:ParametricNode{T},NT<:NodeTangent{T,N},PARAM<:AbstractMatrix{T}}
+    d_constants = extract_gradient(gradient.tree, get_tree(ex))
+    d_params = gradient.metadata._data.parameters[:]
+    return vcat(d_constants, d_params)  # Same shape as `get_constants`
+end
 
 function Base.convert(::Type{Node}, ex::ParametricExpression{T}) where {T}
     num_params = UInt16(size(ex.metadata.parameters, 1))

test/test_parametric_expression.jl

@@ -316,4 +316,8 @@ end
     @test grad.tree.gradient ≈ true_grad[3]
     # Gradient w.r.t. the parameters:
     @test grad.metadata._data.parameters ≈ true_grad[2]
+
+    # Gradient extractor
+    @test extract_gradient(grad, ex) ≈ vcat(true_grad[3], true_grad[2][:])
+    @test axes(extract_gradient(grad, ex)) == axes(first(get_constants(ex)))
 end