diff --git a/ext/ArrayInterfaceChainRulesCoreExt.jl b/ext/ArrayInterfaceChainRulesCoreExt.jl index 6cf4c406..2ef97b5a 100644 --- a/ext/ArrayInterfaceChainRulesCoreExt.jl +++ b/ext/ArrayInterfaceChainRulesCoreExt.jl @@ -5,13 +5,13 @@ import ChainRulesCore import ChainRulesCore: unthunk, NoTangent, ZeroTangent, ProjectTo, @thunk function ChainRulesCore.rrule(::typeof(ArrayInterface.restructure), target, src) - projectT = ProjectTo(target) + projectS = ProjectTo(src) function restructure_pullback(dt) dt = unthunk(dt) f̄ = NoTangent() t̄ = ZeroTangent() - s̄ = @thunk(projectT(ArrayInterface.restructure(src, dt))) + s̄ = @thunk(projectS(ArrayInterface.restructure(src, dt))) f̄, t̄, s̄ end diff --git a/test/chainrules.jl b/test/chainrules.jl index 759a55be..fa56e8e2 100644 --- a/test/chainrules.jl +++ b/test/chainrules.jl @@ -11,3 +11,4 @@ b = zeros(length(arr)) ChainRulesTestUtils.test_rrule(ArrayInterface.restructure, arr, b) ChainRulesTestUtils.test_rrule(ArrayInterface.restructure, b, arr) +ChainRulesTestUtils.test_rrule(ArrayInterface.restructure, zeros(2, 2), zeros(4))