Apply suggestions from code review

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

Author: gaurav-arya

docs/src/implementations.md

@@ -20,10 +20,9 @@ To define a new FFT implementation in your own module, you should
 * Define a new method `AbstractFFTs.plan_fft(x, region; kws...)` that returns a `MyPlan` for at least some types of
   `x` and some set of dimensions `region`.   The `region` (or a copy thereof) should be accessible via `fftdims(p::MyPlan)` (which defaults to `p.region`).
 
-* Define a method of `LinearAlgebra.mul!(y, p::MyPlan, x)` (or `A_mul_B!(y, p::MyPlan, x)` on Julia prior to
-  0.7.0-DEV.3204) that computes the transform `p` of `x` and stores the result in `y`.
+* Define a method of `LinearAlgebra.mul!(y, p::MyPlan, x)` that computes the transform `p` of `x` and stores the result in `y`.
 
-* Define a method of `*(p::MyPlan, x)`, which can simply call your `mul!` (or `A_mul_B!`) method.
+* Define a method of `*(p::MyPlan, x)`, which can simply call your `mul!` method.
   This is not defined generically in this package due to subtleties that arise for in-place and real-input FFTs.
 
 * If the inverse transform is implemented, you should also define `plan_inv(p::MyPlan)`, which should construct the
@@ -33,7 +32,7 @@ To define a new FFT implementation in your own module, you should
 
 * You can also define similar methods of `plan_rfft` and `plan_brfft` for real-input FFTs.
 
-* To enable automatic computation of adjoint plans via [`Base.adjoint`](@ref) (used in rules for reverse differentiation), define the trait `AbstractFFTs.ProjectionStyle(::MyPlan)`, which can take values:
+* To enable automatic computation of adjoint plans via [`Base.adjoint`](@ref) (used in rules for reverse-mode differentiation), define the trait `AbstractFFTs.ProjectionStyle(::MyPlan)`, which can take values:
     * `AbstractFFTs.NoProjectionStyle()`,
     * `AbstractFFTs.RealProjectionStyle()`, for plans which halve one of the output's dimensions analogously to [`rfft`](@ref),
     * `AbstractFFTs.RealInverseProjectionStyle(d::Int)`, for plans which expect an input with a halved dimension analogously to [`irfft`](@ref), where `d` is the original length of the dimension.

src/chainrules.jl

@@ -170,20 +170,19 @@ end
 
 function ChainRulesCore.frule((_, ΔP, Δx), ::typeof(*), P::ScaledPlan, x::AbstractArray) 
     y = P * x 
-    Δy = P * Δx + ΔP.scale / P.scale * y
+    Δy = P * Δx .+ (ΔP.scale / P.scale) .* y
     return y, Δy
 end
 function ChainRulesCore.rrule(::typeof(*), P::ScaledPlan, x::AbstractArray)
     y = P * x
     project_x = ChainRulesCore.ProjectTo(x)
-    project_scale = ChainRulesCore.ProjectTo(P.scale)
     Pt = P'
     scale = P.scale
-    function mul_plan_pullback(ȳ)
+    function mul_scaledplan_pullback(ȳ)
         x̄ = ChainRulesCore.@thunk(project_x(Pt * ȳ))
-        scale_tangent = ChainRulesCore.@thunk(project_scale(sum(conj(y) .* ȳ) / conj(scale)))
+        scale_tangent = ChainRulesCore.@thunk(dot(y, ȳ) / conj(scale))
         plan_tangent = ChainRulesCore.Tangent{typeof(P)}(;p=ChainRulesCore.NoTangent(), scale=scale_tangent)
         return ChainRulesCore.NoTangent(), plan_tangent, x̄ 
     end
-    return y, mul_plan_pullback
+    return y, mul_scaledplan_pullback
 end

test/runtests.jl

@@ -206,7 +206,7 @@ end
             y = randn(size(x))
             for dims in unique((1, 1:N, N))
                 P = plan_fft(x, dims)
-                @test AbstractFFTs.output_size(P) == size(x)
+                @test @inferred(AbstractFFTs.output_size(P)) == size(x)
                 @test AbstractFFTs.output_size(P') == size(x)
                 Pinv = plan_ifft(x)
                 @test AbstractFFTs.output_size(Pinv) == size(x)
@@ -222,7 +222,7 @@ end
                 Px_sz = size(P * x)
                 @test AbstractFFTs.output_size(P) == Px_sz 
                 @test AbstractFFTs.output_size(P') == size(x) 
-                y = randn(Px_sz) .+ randn(Px_sz) * im
+                y = randn(Complex{Float64}, Px_sz)
                 Pinv = plan_irfft(y, size(x)[first(dims)], dims)
                 @test AbstractFFTs.output_size(Pinv) == size(Pinv * y)
                 @test AbstractFFTs.output_size(Pinv') == size(y)
@@ -256,9 +256,7 @@ end
             N = ndims(x)
             for dims in unique((1, 1:N, N))
                 P = plan_rfft(x, dims)        
-                y_real = randn(size(P * x))
-                y_imag = randn(size(P * x))
-                y = y_real .+ y_imag .* im 
+                y = randn(Complex{Float64}, size(P * x))
                 @test (P')' * x == P * x
                 @test size(P') == AbstractFFTs.output_size(P) 
                 @test dot(y_real, real.(P * x)) + dot(y_imag, imag.(P * x)) ≈ dot(P' * y, x)