Add rules and tests for ScaledPlan

Author: gaurav-arya

Project.toml

@@ -12,9 +12,10 @@ julia = "^1.0"
 
 [extras]
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["ChainRulesTestUtils", "Random", "Test", "Unitful"]
+test = ["ChainRulesTestUtils", "FiniteDifferences", "Random", "Test", "Unitful"]

src/chainrules.jl

@@ -167,3 +167,23 @@ function ChainRulesCore.rrule(::typeof(*), P::Plan, x::AbstractArray)
     end
     return y, mul_plan_pullback
 end
+
+function ChainRulesCore.frule((_, ΔP, Δx), ::typeof(*), P::ScaledPlan, x::AbstractArray) 
+    y = P * x 
+    Δ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(ȳ)
+        x̄ = ChainRulesCore.@thunk(project_x(Pt * ȳ))
+        scale_tangent = ChainRulesCore.@thunk(project_scale(sum(conj(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
+end

test/runtests.jl

@@ -1,9 +1,10 @@
 # This file contains code that was formerly part of Julia. License is MIT: https://julialang.org/license
 
 using AbstractFFTs
-using AbstractFFTs: Plan
+using AbstractFFTs: Plan, ScaledPlan
 using ChainRulesTestUtils
-using ChainRulesCore: NoTangent
+using FiniteDifferences
+import ChainRulesCore
 
 using LinearAlgebra
 using Random
@@ -293,9 +294,21 @@ end
     end
 
     @testset "fft" begin
-        for x in (randn(2), randn(2, 3), randn(3, 4, 5))
-            N = ndims(x)
-            complex_x = complex.(x)
+        # Overloads to allow ChainRulesTestUtils to test rules w.r.t. ScaledPlan's. See https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/256
+        InnerPlan = Union{TestPlan, InverseTestPlan, TestRPlan, InverseTestRPlan}
+        function FiniteDifferences.to_vec(x::InnerPlan)
+            function FFTPlan_from_vec(x_vec::Vector)
+                return x
+            end
+            return Bool[], FFTPlan_from_vec
+        end
+        ChainRulesTestUtils.test_approx(::ChainRulesCore.AbstractZero, x::InnerPlan, msg=""; kwargs...) = true
+        ChainRulesTestUtils.rand_tangent(::AbstractRNG, x::InnerPlan) = ChainRulesCore.NoTangent()
+
+        for x_shape in ((2,), (2, 3), (3, 4, 5))
+            N = length(x_shape)
+            x = randn(x_shape)
+            complex_x = x + randn(x_shape) * im
             for dims in unique((1, 1:N, N))
                 # fft, ifft, bfft
                 for f in (fft, ifft, bfft)
@@ -305,17 +318,17 @@ end
                     test_rrule(f, complex_x, dims)
                 end
                 for pf in (plan_fft, plan_ifft, plan_bfft) 
-                    test_frule(*, pf(x, dims) ⊢ NoTangent(), x)
-                    test_rrule(*, pf(x, dims) ⊢ NoTangent(), x)
-                    test_frule(*, pf(complex_x, dims) ⊢ NoTangent(), complex_x)
-                    test_rrule(*, pf(complex_x, dims) ⊢ NoTangent(), complex_x)
+                    test_frule(*, pf(x, dims), x)
+                    test_rrule(*, pf(x, dims), x)
+                    test_frule(*, pf(complex_x, dims), complex_x)
+                    test_rrule(*, pf(complex_x, dims), complex_x)
                 end
 
                 # rfft 
                 test_frule(rfft, x, dims)
                 test_rrule(rfft, x, dims)
-                test_frule(*, plan_rfft(x, dims) ⊢ NoTangent(), x)
-                test_rrule(*, plan_rfft(x, dims) ⊢ NoTangent(), x)
+                test_frule(*, plan_rfft(x, dims), x)
+                test_rrule(*, plan_rfft(x, dims), x)
 
                 # irfft, brfft
                 for f in (irfft, brfft)
@@ -328,8 +341,8 @@ end
                 end
                 for pf in (plan_irfft, plan_brfft) 
                     for d in (2 * size(x, first(dims)) - 1, 2 * size(x, first(dims)) - 2)
-                        test_frule(*, pf(complex_x, d, dims) ⊢ NoTangent(), complex_x)
-                        test_rrule(*, pf(complex_x, d, dims) ⊢ NoTangent(), complex_x)
+                        test_frule(*, pf(complex_x, d, dims), complex_x)
+                        test_rrule(*, pf(complex_x, d, dims), complex_x)
                     end
                 end
             end