Merge pull request #119 from SymbolicML/compathelper/new_version/2025-01-05-01-23-24-852-00256837307

CompatHelper: bump compat for Zygote in [weakdeps] to 0.7, (keep existing compat)

Author: MilesCranmer

Project.toml

@@ -38,7 +38,7 @@ Optim = "0.19, 1"
 PrecompileTools = "1"
 Reexport = "1"
 SymbolicUtils = "0.19, ^1.0.5, 2, 3"
-Zygote = "0.6"
+Zygote = "0.6, 0.7"
 julia = "1.10"
 
 [extras]

src/ChainRules.jl

@@ -7,6 +7,7 @@ using ChainRulesCore:
     ZeroTangent,
     Tangent,
     @thunk,
+    unthunk,
     canonicalize
 using ..OperatorEnumModule: OperatorEnum
 using ..NodeModule: AbstractExpressionNode, with_type_parameters, tree_mapreduce
@@ -52,7 +53,8 @@ struct EvalPullback{N,A,O} <: Function
 end
 
 # TODO: Preferable to use the primal in the pullback somehow
-function (e::EvalPullback)((dY, _))
+function (e::EvalPullback)((thunked_dY, _))
+    dY = unthunk(thunked_dY)
     _, dX_constants_dY, complete = eval_grad_tree_array(
         e.tree, e.X, e.operators; variable=Val(:both)
     )