definitions for adjoints of FFTW functions

[sipposip]

Adjoint definitions and accompanying tests for the most important functions of FFTW (fft and ifft).

Since deriving the adjoint of FFTs is a bit tricky, I double-checked the results with differentiating a simple DFT-implementation which is differentiable with ReverseDiff, and the results are equal down to machine-precision. This is possible since FFT is in fact only a more efficient implementation of DFT, which itself boils down to a simple matrix-multiplication.

I add this comparison as comment here for reference:

using ReverseDiff
using FFTW
using Test
using Zygote
using FillArrays: Fill


function dft(x::AbstractArray)
    """discrete fourier transform"""
    N = size(x)[1]
    out = Array{Any}(undef,N)
    for k in 0:N-1
        out[k+1] = sum([x[n+1]*exp(-2*im*π*k*n/N) for n in 0:N-1])
    end
    return out
end

function idft(x::AbstractArray)
    """discrete inverse fourier transform"""
    N = size(x)[1]
    out = Array{Any}(undef,N)
    for n in 0:N-1
        out[n+1] = 1/N*sum([x[k+1]*exp(2*im*π*k*n/N) for k in 0:N-1])
    end
    return out
end


x = randn(100)
@test dft(x) ≈ FFTW.fft(x)
@test idft(x) ≈ FFTW.ifft(x)


# FFTW functions do not work with FillArray. To make it work with FillArrays
# as well, overload the functions
FFTW.fft(x::Fill) = FFTW.fft(collect(x))
FFTW.ifft(x::Fill) = FFTW.ifft(collect(x))

# the gradient of an FFT with respect to its input is the reverse FFT of the
# gradient of its inputs.
Zygote.@adjoint function FFTW.fft(xs)
    return FFTW.fft(xs), function(Δ)
        N = length(Δ)
        return (N * FFTW.ifft(Δ),)
    end
end

Zygote.@adjoint function FFTW.ifft(xs)
    return FFTW.ifft(xs), function(Δ)
        N = length(Δ)
        return (1/N* FFTW.fft(Δ),)
    end
end


@test real(Zygote.gradient((x)->sum(abs.(FFTW.fft(x))),x)[1]) ≈ ReverseDiff.gradient((x)->sum(abs.(dft(x))),x)
@test real(Zygote.gradient((x)->sum(abs.(FFTW.ifft(x))),x)[1]) ≈ ReverseDiff.gradient((x)->sum(abs.(idft(x))),x)

[MikeInnes]

Thanks! Just a couple comments but this is great to have.

[MikeInnes]

You can use dims... to make this more compact.

[sipposip]

thanks for the tip, I changed to dims... now

[MikeInnes]

Please use the package manager to add FFTW, so that the formatting is right and so that the manifest gets updated too.

[sipposip]

fixed now via package manager

[MikeInnes]

I think this file is short enough that it could just go in array.jl, after linear algebra.

[sipposip]

I moved it to array.jl now

[MikeInnes]

I suspect this is unnecessary, since numbers act like single-element collections anyway for things like prod.

[sipposip]

the exception for int was necessary because size(xs)[collect(1)] does not work, but I now changed to collect(size(xs))[dims] which works and seems more elegant to me and does not need different handling for ints

[MikeInnes]

This comment isn't clear to me; you seem to be handling multiple dimensions below. If only single integers work, best to add a type restriction here.

[sipposip]

this comment was from an old version and I forgot to remove it, sorry for the confusion. It indeed works also with multiple dimensions. I removed the comment now

[MikeInnes]

It would be a bit simpler to use gradcheck here instead. That will check against numerical gradients, which is handy for correctness.

[sipposip]

I added 4 gradchecks as well now. I think the old test still make sense since they check whether the FFTs along individual dims work and produce the same results as when doing all dimensions in one go. However if you think that the old tests are mainly bloating gradcheck.jl I can also ommit them

[MikeInnes]

The indentation is off here, just needs a quick fix.

[jekbradbury]

nit: indent these few lines?

[sipposip]

I (hopefully) fixed the indentation now.

[MikeInnes]

Alright, great, this looks good.

There's a merge conflict due to the manifest; it'd be great if you could resolve that, and even better if you can use a newer Julia version to minimise the number of changes while you're at it. Other than that I think we're good to go.

[sipposip]

OK. I am very new both to Julia and collaborating on github, but I will do my best to resolve the conflicts

[sipposip]

I now switched to julia 1.1.1 and resolved the conflicts. However, for some reason my manifest has more entries than the current one in master (AbstractFFTs, Conda, JSON, Reexport, VersionParsing) in addition to FFTW that I added via the package manager. Should I remove them?

[MikeInnes]

No, that's ok, they are just dependencies of FFTW that got pulled in.

Unfortunately the tests are now failing; looks like you just have a syntax error in the tests you added.

[sipposip]

I now fixed the tests, and they pass on my local installation. I also fixed the testset for hcat which had failed as well, and some other fixes in other testsets (forgotten @test)

[sipposip]

The travis CI has finished, but this does not seem to be reported to github (when I check https://travis-ci.org/FluxML/Zygote.jl/builds/543743653?utm_source=github_status&utm_medium=notification all tests have passed, but here it is still pending.) Is there any way to re-trigger the travis-ci build?

[ChrisRackauckas]

Can't you just depend on AbstractFFTs.jl and write the dispatches using the high level functions?

[MikeInnes]

👍 Thanks!

[ChrisRackauckas]

But can't this be done with just AbstractFFTs?

[MikeInnes]

Probably, but fine to do that as a follow up. The main issue is testing but we can have FFTW as a test dependency.