Skip to content

Commit 840bdf8

Browse files
Precompute plan transpose for rrule
Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
1 parent 31c3eac commit 840bdf8

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

src/chainrules.jl

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -160,9 +160,10 @@ end
160160
function ChainRulesCore.rrule(::typeof(*), P::Plan, x::AbstractArray)
161161
y = P * x
162162
project_x = ChainRulesCore.ProjectTo(x)
163-
function fft_pullback(ȳ)
164-
= project_x(P' * ȳ)
163+
Pt = P'
164+
function mul_plan_pullback(ȳ)
165+
= project_x(Pt * ȳ)
165166
return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), x̄
166167
end
167-
return y, fft_pullback
168+
return y, mul_plan_pullback
168169
end

0 commit comments

Comments
 (0)