fix: unthunk ChainRules pullback outputs (#674)

Author: gdalle

DifferentiationInterface/ext/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

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