Increasing robustness of adjoint plan optimizations

src/definitions.jl

@@ -683,10 +683,11 @@
     pinv = inv(p)
     # Optimization: when pinv is a ScaledPlan, check if we can avoid a loop over x.
     # Even if not, ensure that we do only one pass by combining the normalization with the plan.
+    RT = AbstractArray{<:T} # canonicalize eltype of returned array to be extra safe about type stability  
     if pinv isa ScaledPlan && pinv.scale == N
-        return pinv.p * x
+        return convert(RT, pinv.p * x)
     else
-        return (1/N * pinv) * x
+        return convert(RT, (inv(N) * pinv) * x)
     end
 end
 
@@ -698,7 +699,7 @@
     pinv = inv(p)
     n = size(pinv, halfdim)
     # Optimization: when pinv is a ScaledPlan, fuse the scaling into our map to ensure we do not loop over x twice.
-    scale = pinv isa ScaledPlan ? pinv.scale / 2N : 1 / 2N
+    scale = pinv isa ScaledPlan ? pinv.scale / 2N : one(pinv.scale) / 2N
     twoscale = 2 * scale 
     unscaled_pinv = pinv isa ScaledPlan ? pinv.p : pinv 
     y = map(x, CartesianIndices(x)) do xj, j
@@ -721,7 +722,7 @@
     pinv = inv(p)
     d = size(pinv, halfdim)
     # Optimization: when pinv is a ScaledPlan, fuse the scaling into our map to ensure we do not loop over x twice.
-    scale = pinv isa ScaledPlan ? pinv.scale / N : 1 / N
+    scale = pinv isa ScaledPlan ? pinv.scale / N : one(pinv.scale) / N
     twoscale = 2 * scale 
     unscaled_pinv = pinv isa ScaledPlan ? pinv.p : pinv 
     y = unscaled_pinv * x

test/runtests.jl

@@ -145,6 +145,21 @@ end
     end
 end
 
+@testset "Adjoint plan on single-precision" begin
+    # fft
+    p = plan_fft(zeros(ComplexF32, 3))
+    u = rand(ComplexF32, 3)
+    @test eltype(p' * (p * u)) == eltype(u)
+    # rfft
+    p = plan_rfft(zeros(Float32, 3))
+    u = rand(Float32, 3)
+    @test eltype(p' * (p * u)) == eltype(u)
+    # brfft
+    p = plan_brfft(zeros(ComplexF32, 3), 5)
+    u = rand(ComplexF32, 3)
+    @test eltype(p' * (p * u)) == eltype(u)
+end
+
 @testset "ChainRules" begin
     @testset "shift functions" begin
         for x in (randn(3), randn(3, 4), randn(3, 4, 5))