definitions for adjoints of FFTW functions

Manifest.toml

@@ -1,5 +1,11 @@
 # This file is machine-generated - editing it directly is not advised
 
+[[AbstractFFTs]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "380e36c66edfa099cd90116b24c1ce8cafccac40"
+uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+version = "0.4.1"
+
 [[Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
@@ -33,6 +39,12 @@ git-tree-sha1 = "84aa74986c5b9b898b0d1acaf3258741ee64754f"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
 version = "2.1.0"
 
+[[Conda]]
+deps = ["Compat", "JSON", "VersionParsing"]
+git-tree-sha1 = "b625d802587c2150c279a40a646fba63f9bd8187"
+uuid = "8f4d0f93-b110-5947-807f-2305c1781a2d"
+version = "1.2.0"
+
 [[Crayons]]
 deps = ["Test"]
 git-tree-sha1 = "f621b8ef51fd2004c7cf157ea47f027fdeac5523"
@@ -69,6 +81,12 @@ version = "0.0.10"
 deps = ["Random", "Serialization", "Sockets"]
 uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 
+[[FFTW]]
+deps = ["AbstractFFTs", "BinaryProvider", "Compat", "Conda", "Libdl", "LinearAlgebra", "Reexport", "Test"]
+git-tree-sha1 = "29cda58afbf62f35b1a094882ad6c745a47b2eaa"
+uuid = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+version = "0.2.4"
+
 [[FillArrays]]
 deps = ["LinearAlgebra", "Random", "SparseArrays", "Test"]
 git-tree-sha1 = "9ab8f76758cbabba8d7f103c51dce7f73fcf8e92"
@@ -93,6 +111,12 @@ version = "0.2.1"
 deps = ["Markdown"]
 uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 
+[[JSON]]
+deps = ["Dates", "Distributed", "Mmap", "Sockets", "Test", "Unicode"]
+git-tree-sha1 = "1f7a25b53ec67f5e9422f1f551ee216503f4a0fa"
+uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+version = "0.20.0"
+
 [[LibGit2]]
 uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
 
@@ -153,6 +177,12 @@ uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 deps = ["Serialization"]
 uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
+[[Reexport]]
+deps = ["Pkg"]
+git-tree-sha1 = "7b1d07f411bc8ddb7977ec7f377b97b158514fe0"
+uuid = "189a3867-3050-52da-a836-e630ba90ab69"
+version = "0.2.0"
+
 [[Requires]]
 deps = ["Test"]
 git-tree-sha1 = "f6fbf4ba64d295e146e49e021207993b6b48c7d1"
@@ -219,3 +249,9 @@ uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [[Unicode]]
 uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
+
+[[VersionParsing]]
+deps = ["Compat"]
+git-tree-sha1 = "c9d5aa108588b978bd859554660c8a5c4f2f7669"
+uuid = "81def892-9a0e-5fdd-b105-ffc91e053289"
+version = "1.1.3"

Project.toml

@@ -4,6 +4,7 @@ version = "0.3.2"
 
 [deps]
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 IRTools = "7869d1d1-7146-5819-86e3-90919afe41df"

src/lib/array.jl

@@ -1,4 +1,4 @@
-using FillArrays
+using FillArrays, FFTW
 using FillArrays: AbstractFill, getindex_value
 using Base.Broadcast: broadcasted, broadcast_shape
 
@@ -429,6 +429,53 @@ end
 @adjoint -(A::AbstractArray, B::AbstractArray) = A - B, Δ->(Δ, -Δ)
 @adjoint -(A::AbstractArray) = -A, Δ->(-Δ,)
 
+# FFTW
+# ===================
+
+# FFTW functions do not work with FillArrays, which are needed
+# for some functionality of Zygote. To make it work with FillArrays
+# as well, overload the relevant functions
+FFTW.fft(x::Fill, dims...) = FFTW.fft(collect(x), dims...)
+FFTW.ifft(x::Fill, dims...) = FFTW.ifft(collect(x), dims...)
+
+
+# the adjoint jacobian of an FFT with respect to its input is the reverse FFT of the
+# gradient of its inputs, but with different normalization factor
+@adjoint function FFTW.fft(xs)
+  return FFTW.fft(xs), function(Δ)
+    N = length(xs)
+    return (N * FFTW.ifft(Δ),)
+  end
+end
+
+@adjoint function FFTW.ifft(xs)
+  return FFTW.ifft(xs), function(Δ)
+    N = length(xs)
+    return (1/N* FFTW.fft(Δ),)
+  end
+end
+
+@adjoint function FFTW.fft(xs, dims)
+  return FFTW.fft(xs, dims), function(Δ)
+    # dims can be int, array or tuple,
+    # convert to collection for use as index
+    dims = collect(dims)
+    # we need to multiply by all dimensions that we FFT over
+    N = prod(collect(size(xs))[dims])
+    return (N * FFTW.ifft(Δ, dims), nothing)
+  end
+end
+
+@adjoint function FFTW.ifft(xs,dims)
+  return FFTW.ifft(xs, dims), function(Δ)
+    # dims can be int, array or tuple,
+    # convert to collection for use as index
+    dims = collect(dims)
+    # we need to divide by all dimensions that we FFT over
+    N = prod(collect(size(xs))[dims])
+    return (1/N * FFTW.fft(Δ, dims),nothing)
+  end
+end
 
 # FillArray functionality
 # =======================

test/gradcheck.jl

@@ -1,4 +1,4 @@
-using Zygote, NNlib, Test, Random, LinearAlgebra, Statistics, FillArrays
+using Zygote, NNlib, Test, Random, LinearAlgebra, Statistics, FillArrays, FFTW
 using Zygote: gradient
 using NNlib: conv, ∇conv_data, depthwiseconv
 using Base.Broadcast: broadcast_shape
@@ -563,17 +563,38 @@ using Zygote: Buffer
 end
 
 @testset "FillArrays" begin
-  gradcheck(x->sum(Fill(x[], (2, 2))), [0.1])
+  @test gradcheck(x->sum(Fill(x[], (2, 2))), [0.1])
   @test first(Zygote.gradient(sz->sum(Ones(sz)), 6)) === nothing
   @test first(Zygote.gradient(sz->sum(Zeros(sz)), 6)) === nothing
 end
 
 @testset "AbstractArray Addition / Subtraction / Negation" begin
   rng, M, N, P = MersenneTwister(123567), 3, 7, 11
   A, B = randn(rng, M, N, P), randn(rng, M, N, P)
-  gradtest(+, A, B)
-  gradtest(-, A, B)
-  gradtest(-, A)
+  @test gradtest(+, A, B)
+  @test gradtest(-, A, B)
+  @test gradtest(-, A)
+end
+
+@testset "FFTW" begin
+  x=[-0.353213 -0.789656 -0.270151; -0.95719 -1.27933 0.223982]
+  # gradient of ifft(rfft) must be 1
+  @test gradient((x)->real(ifft(fft(x))[1]),x)[1][1] == 1.0+0.0im
+  @test gradient((x)->real(fft(ifft(x))[1]),x)[1][1] == 1.0+0.0im
+
+  # check ffts for individual dimensions
+  @test gradient((x)->sum(abs.(FFTW.fft(x))),x)[1] ≈ gradient((x)->sum(abs.(FFTW.fft(FFTW.fft(x,1),2))),x)[1]
+  @test gradient((x)->abs(sum((FFTW.fft(x)))),x)[1] ≈ gradient((x)->abs(sum(FFTW.fft(FFTW.fft(x,1),2))),x)[1]
+  @test gradient((x, dims)->sum(abs.(FFTW.fft(x,dims))),x,(1,2))[1] ≈ gradient((x)->sum(abs.(FFTW.fft(x))),x)[1]
+  @test gradient((x)->sum(abs.(FFTW.fft(x,(1,2)))),x)[1] ≈ gradient((x)->sum(abs.(FFTW.fft(FFTW.fft(x,1),2))),x)[1]
+  @test gradient((x, dims)->sum(abs.(FFTW.ifft(x,dims))),x,(1,2))[1] ≈ gradient((x)->sum(abs.(FFTW.ifft(x))),x)[1]
+  @test gradient((x)->sum(abs.(FFTW.ifft(x,(1,2)))),x)[1] ≈ gradient((x)->sum(abs.(FFTW.ifft(FFTW.ifft(x,1),2))),x)[1]
+
+  @test gradcheck(x->sum(abs.(FFTW.fft(x))), x)
+  @test gradcheck(x->sum(abs.(FFTW.ifft(x))), x)
+  @test gradcheck(x->sum(abs.(FFTW.fft(x, 1))), x)
+  @test gradcheck(x->sum(abs.(FFTW.ifft(x, 1))), x)
+
 end
 
 @testset "FillArrays" begin