Remove extra unthunk_tangent

FluxML · Jan 3, 2025 · 2d21e2e · 2d21e2e
1 parent d7ecca6
commit 2d21e2e
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/lib/lib.jl b/src/lib/lib.jl
@@ -37,9 +37,12 @@ function accum(x::RefValue, y::RefValue)
   return x
 end
 
-accum(x, y::AbstractThunk) = @thunk(accum(x, unthunk_tangent(y)))
-accum(x::AbstractThunk, y) = @thunk(accum(unthunk_tangent(x), y))
-accum(x::AbstractThunk, y::AbstractThunk) = @thunk(accum(unthunk_tangent(x), unthunk_tangent(y)))
+accum(x::NamedTuple, y::ChainRulesCore.Tangent) = accum(x, wrap_chainrules_output(y))
+accum(x::ChainRulesCore.Tangent, y::NamedTuple) = accum(wrap_chainrules_output(x), y)
+
+accum(x, y::AbstractThunk) = @thunk(accum(x, unthunk(y)))
+accum(x::AbstractThunk, y) = @thunk(accum(unthunk(x), y))
+accum(x::AbstractThunk, y::AbstractThunk) = @thunk(accum(unthunk(x), unthunk(y)))
 
 # Core functions
 @_adjoint_keepthunks deepcopy(x) = deepcopy(x), ȳ -> (ȳ,)