Closed pzimbrod closed 2 years ago
This problem is mentioned in FluxML/Zygote.jl#899 and subsequently JuliaMath/FFTW.jl#182. In some implementations of FFTW.jl
the region
field of the struct plan is missing.
It would be plausible that the problem can be fixed by switching to the unscaled inverse transform and do the scaling manually afterwards.
Apparently, you can use plan_brfft
instead for the inverse and scale afterwards:
fib = plan_brfft(rfft(n), length(n));
# Check for same results
irfft(f * n, length(n)) ≈ fib * (f * n) ./ length(n)
true
# Does Zygote run now?
jacobian(x -> fib * (f * x) ./ length(x), n)
([0.51 0.0 … 0.0 0.0; 0.0 0.51 … 0.0 3.996891658508158e-19; … ; 0.0 0.0 … 0.51 0.0; 0.0 8.673617379884035e-19 … 0.0 0.51],)
When taking the jacobian of a sample FFT-Pipeline,
Zygote
complains about missing fields in the corresponding FFT Plan:However, doing the FFT in-place, the error vanishes:
This is an issue since we're doing the FFT and its inverse a lot of times, so the speed gain from using pre-planned FFTs shoud be quite considerable.