ForwardDiff and Zygote cannot automatically differentiate (AD) function from C^n to R that uses FFT

Specific Domains Numerics
4

I think the issue is that ForwardDiff does not work with FFTW, and Zygote’s hessian function is just hessian(f, x::AbstractArray) = forward_jacobian(x -> gradient(f, x)[1], x)[2] i.e. it calculates the gradient, and then uses ForwardDiff. to differentiate again.

julia> gradient(x->sum(abs2, fftshift(fft(x))), rand(3) .+ im .* randn.())
(ComplexF64[1.495784118230028 - 13.219849172255007im, 4.629297683031974 + 5.3532713394798055im, 2.672875552797116 - 0.3633324783983167im],)

julia> Zygote.hessian(x->sum(abs2, fft(x)), rand(10))
ERROR: type ForwardDiff.Dual{Nothing, Float64, 10} not supported

You might be able to cook something up using say this jacobian definition, so that both derivatives are taken using Zygote.

There might also be a more generic FFT somewhere which does accept dual numbers, i.e. not FFTW. Or as @Tamas_Papp suggests it might be possible to add a rule for this somewhere — Zygote already has rules which write the reverse gradient in terms of other ffts. ForwardDiff does not use ChainRules, but I don’t see why you couldn’t write an explicit method AbstractFFTs.fft(::AbstractArray{Dual{... to handle this.

3 Likes