Implement FFT chain rules using adjoint plans

Author: gaurav-arya

src/chainrules.jl

@@ -1,123 +1,43 @@
-# ffts
-function ChainRulesCore.frule((_, Δx, _), ::typeof(fft), x::AbstractArray, dims)
-    y = fft(x, dims)
-    Δy = fft(Δx, dims)
-    return y, Δy
-end
-function ChainRulesCore.rrule(::typeof(fft), x::AbstractArray, dims)
-    y = fft(x, dims)
-    project_x = ChainRulesCore.ProjectTo(x)
-    function fft_pullback(ȳ)
-        x̄ = project_x(bfft(ChainRulesCore.unthunk(ȳ), dims))
-        return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent()
-    end
-    return y, fft_pullback
-end
-
-function ChainRulesCore.frule((_, Δx, _), ::typeof(rfft), x::AbstractArray{<:Real}, dims)
-    y = rfft(x, dims)
-    Δy = rfft(Δx, dims)
-    return y, Δy
-end
-function ChainRulesCore.rrule(::typeof(rfft), x::AbstractArray{<:Real}, dims)
-    y = rfft(x, dims)
-
-    # compute scaling factors
-    halfdim = first(dims)
-    d = size(x, halfdim)
-    n = size(y, halfdim)
-    scale = reshape(
-        [i == 1 || (i == n && 2 * (i - 1) == d) ? 1 : 2 for i in 1:n],
-        ntuple(i -> i == first(dims) ? n : 1, Val(ndims(x))),
-    )
-
-    project_x = ChainRulesCore.ProjectTo(x)
-    function rfft_pullback(ȳ)
-        x̄ = project_x(brfft(ChainRulesCore.unthunk(ȳ) ./ scale, d, dims))
-        return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent()
-    end
-    return y, rfft_pullback
-end
-
-function ChainRulesCore.frule((_, Δx, _), ::typeof(ifft), x::AbstractArray, dims)
-    y = ifft(x, dims)
-    Δy = ifft(Δx, dims)
-    return y, Δy
-end
-function ChainRulesCore.rrule(::typeof(ifft), x::AbstractArray, dims)
-    y = ifft(x, dims)
-    invN = normalization(y, dims)
-    project_x = ChainRulesCore.ProjectTo(x)
-    function ifft_pullback(ȳ)
-        x̄ = project_x(invN .* fft(ChainRulesCore.unthunk(ȳ), dims))
-        return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent()
-    end
-    return y, ifft_pullback
-end
-
-function ChainRulesCore.frule((_, Δx, _, _), ::typeof(irfft), x::AbstractArray, d::Int, dims)
-    y = irfft(x, d, dims)
-    Δy = irfft(Δx, d, dims)
-    return y, Δy
-end
-function ChainRulesCore.rrule(::typeof(irfft), x::AbstractArray, d::Int, dims)
-    y = irfft(x, d, dims)
-
-    # compute scaling factors
-    halfdim = first(dims)
-    n = size(x, halfdim)
-    invN = normalization(y, dims)
-    twoinvN = 2 * invN
-    scale = reshape(
-        [i == 1 || (i == n && 2 * (i - 1) == d) ? invN : twoinvN for i in 1:n],
-        ntuple(i -> i == first(dims) ? n : 1, Val(ndims(x))),
-    )
-
-    project_x = ChainRulesCore.ProjectTo(x)
-    function irfft_pullback(ȳ)
-        x̄ = project_x(scale .* rfft(real.(ChainRulesCore.unthunk(ȳ)), dims))
-        return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent()
+for f in (:fft, :bfft, :ifft, :rfft)
+    pf = Symbol("plan_", f)
+    @eval begin
+        function ChainRulesCore.frule((_, Δx, _), ::typeof($f), x::AbstractArray, dims) 
+            y = $f(x, dims)
+            Δy = $f(Δx, dims)
+            return y, Δy
+        end
+        function ChainRulesCore.rrule(::typeof($f), x::T, dims) where {T<:AbstractArray}
+            y = $f(x, dims)
+            project_x = ChainRulesCore.ProjectTo(x)
+            ax = axes(x)
+            function fft_pullback(ȳ)
+                x̄ = project_x($pf(similar(T, ax), dims)' * ChainRulesCore.unthunk(ȳ))
+                return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent()
+            end
+            return y, fft_pullback
+        end
     end
-    return y, irfft_pullback
 end
 
-function ChainRulesCore.frule((_, Δx, _), ::typeof(bfft), x::AbstractArray, dims)
-    y = bfft(x, dims)
-    Δy = bfft(Δx, dims)
-    return y, Δy
-end
-function ChainRulesCore.rrule(::typeof(bfft), x::AbstractArray, dims)
-    y = bfft(x, dims)
-    project_x = ChainRulesCore.ProjectTo(x)
-    function bfft_pullback(ȳ)
-        x̄ = project_x(fft(ChainRulesCore.unthunk(ȳ), dims))
-        return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent()
+for f in (:brfft, :irfft)
+    pf = Symbol("plan_", f)
+    @eval begin
+        function ChainRulesCore.frule((_, Δx, _), ::typeof($f), x::AbstractArray, d::Int, dims) 
+            y = $f(x, d::Int, dims)
+            Δy = $f(Δx, d::Int, dims)
+            return y, Δy
+        end
+        function ChainRulesCore.rrule(::typeof($f), x::T, d::Int, dims) where {T<:AbstractArray}
+            y = $f(x, d, dims)
+            project_x = ChainRulesCore.ProjectTo(x)
+            ax = axes(x)
+            function fft_pullback(ȳ)
+                x̄ = project_x($pf(similar(T, ax), d, dims)' * ChainRulesCore.unthunk(ȳ))
+                return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent()
+            end
+            return y, fft_pullback
+        end
     end
-    return y, bfft_pullback
-end
-
-function ChainRulesCore.frule((_, Δx, _, _), ::typeof(brfft), x::AbstractArray, d::Int, dims)
-    y = brfft(x, d, dims)
-    Δy = brfft(Δx, d, dims)
-    return y, Δy
-end
-function ChainRulesCore.rrule(::typeof(brfft), x::AbstractArray, d::Int, dims)
-    y = brfft(x, d, dims)
-
-    # compute scaling factors
-    halfdim = first(dims)
-    n = size(x, halfdim)
-    scale = reshape(
-        [i == 1 || (i == n && 2 * (i - 1) == d) ? 1 : 2 for i in 1:n],
-        ntuple(i -> i == first(dims) ? n : 1, Val(ndims(x))),
-    )
-
-    project_x = ChainRulesCore.ProjectTo(x)
-    function brfft_pullback(ȳ)
-        x̄ = project_x(scale .* rfft(real.(ChainRulesCore.unthunk(ȳ)), dims))
-        return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent()
-    end
-    return y, brfft_pullback
 end
 
 # shift functions
@@ -150,3 +70,19 @@ function ChainRulesCore.rrule(::typeof(ifftshift), x::AbstractArray, dims)
     end
     return y, ifftshift_pullback
 end
+
+# plans
+function ChainRulesCore.frule((_, _, Δx), ::typeof(*), P::Plan, x::AbstractArray) 
+    y = P * x 
+    Δy = P * Δx
+    return y, Δy
+end
+function ChainRulesCore.rrule(::typeof(*), P::Plan, x::AbstractArray)
+    y = P * x
+    project_x = ChainRulesCore.ProjectTo(x)
+    function fft_pullback(ȳ)
+        x̄ = project_x(P' * ȳ)
+        return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), x̄
+    end
+    return y, fft_pullback
+end