fix: unthunk ChainRules pullback outputs

JuliaDiff · Jan 6, 2025 · 9f43dc1 · 9f43dc1
1 parent 6604be2
commit 9f43dc1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/...xt/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl b/...xt/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl
@@ -8,7 +8,8 @@ using ChainRulesCore:
     NoTangent,
     RuleConfig,
     frule_via_ad,
-    rrule_via_ad
+    rrule_via_ad,
+    unthunk
 import DifferentiationInterface as DI
 
 ruleconfig(backend::AutoChainRules) = backend.ruleconfig

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/reverse_onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/reverse_onearg.jl
@@ -39,7 +39,7 @@ function DI.value_and_pullback(
     rc = ruleconfig(backend)
     y, pb = rrule_via_ad(rc, f, x, map(DI.unwrap, contexts)...)
     tx = map(ty) do dy
-        pb(dy)[2]
+        unthunk(pb(dy)[2])
     end
     return y, tx
 end
@@ -54,7 +54,7 @@ function DI.value_and_pullback(
 ) where {C}
     (; y, pb) = prep
     tx = map(ty) do dy
-        pb(dy)[2]
+        unthunk(pb(dy)[2])
     end
     return copy(y), tx
 end
@@ -69,7 +69,7 @@ function DI.pullback(
 ) where {C}
     (; pb) = prep
     tx = map(ty) do dy
-        pb(dy)[2]
+        unthunk(pb(dy)[2])
     end
     return tx
 end