# Gradients of FFT

**URL:** <https://discourse.julialang.org/t/gradients-of-fft/70223>\
**Category:** Machine Learning\
**Created:** [October 22, 2021, 3:44pm UTC](https://discourse.julialang.org/t/gradients-of-fft/70223 "2021-10-22T15:44:48Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![maedoc](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maedoc/32/8490_2.png) [@maedoc](https://discourse.julialang.org/u/maedoc)\
**Post date:** [October 22, 2021, 3:44pm UTC](https://discourse.julialang.org/t/gradients-of-fft/70223/1 "2021-10-22T15:44:48Z")

</div>

I tried to expand the DiffEqFlux page’s model to a simple one involving an FFT, but I can’t seem to get it working with any of the `sensalg` options. Here’s what I’ve got,

```julia
using FFTW
using DiffEqSensitivity, OrdinaryDiffEq, Zygote

nlon = 64
nlat = 64
K = (1:(div(nlat,2)+1)) .* (1:nlon)' / (nlat*nlon)

function f(x,p,t)
    k = p[1]
    lc = k * irfft(K .* rfft(x), nlat)
    @. x - x^3/3 + lc
end

p = [1.5];
u0 = randn(nlat, nlon);
prob = ODEProblem(f, u0, (0.0, 10.0), p);
loss(u0,p) = sum(solve(prob,Tsit5(),u0=u0,p=p,saveat=0.1))
du01,dp1 = Zygote.gradient(loss,u0,p)

```

which results in a type error because the FFT package doesn’t support TrackedReals. Doesn’t the AD system recognise what the gradient of an FFT should be and not try to diff through FFTW? This just works in the autograd package I use in Python.

---

<div class="post-metadata">

**Author:** ![oxinabox](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oxinabox/32/206603_2.png) [@oxinabox](https://discourse.julialang.org/u/oxinabox)\
**Post date:** [October 22, 2021, 5:34pm UTC](https://discourse.julialang.org/t/gradients-of-fft/70223/2 "2021-10-22T17:34:40Z")

</div>

Something funky is happening.  
Because you shouldn’t be seeing any `TrackedReal`s.  
That means that either [Tracker.jl](https://github.com/FluxML/Tracker.jl) or [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl) is being used.  
Where as that says you are using `Zygote`.  
Can you post the full error message and stacktrack?

@devmotion has a PR up to add the `rrule` and `frule` definitions to AbstractFFTs.  
[https://github.com/JuliaMath/AbstractFFTs.jl/pull/58](https://github.com/JuliaMath/AbstractFFTs.jl/pull/58)

Once that lands any AD that supports ChainRules (Zygote, Yota, Diffractor, Nabla) should be able to AD `irfft` just fine.

Til then you can try using their branch.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [October 22, 2021, 5:59pm UTC](https://discourse.julialang.org/t/gradients-of-fft/70223/3 "2021-10-22T17:59:08Z")

</div>

Yeah… for reference the current adjoint heuristic is at [https://github.com/SciML/DiffEqSensitivity.jl/blob/v6.60.2/src/concrete\_solve.jl](https://github.com/SciML/DiffEqSensitivity.jl/blob/v6.60.2/src/concrete_solve.jl) . Since your `f` is not in-place, it should be using Zygote for vjps and no tracked values should be seen. So I am confused as well and would like to see that stacktrace.

---

<div class="post-metadata">

**Author:** ![maedoc](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maedoc/32/8490_2.png) [@maedoc](https://discourse.julialang.org/u/maedoc)\
**Post date:** [October 22, 2021, 8:43pm UTC](https://discourse.julialang.org/t/gradients-of-fft/70223/4 "2021-10-22T20:43:56Z")

</div>

Thanks for the replies. I initially tried an in-place form, as in the DiffEqFlux example, but then I thought it might work better out of place. I put the example, traceback and pkg status in this gist

[https://gist.github.com/maedoc/7689780fbb77b9362611eaed6e205da4](https://gist.github.com/maedoc/7689780fbb77b9362611eaed6e205da4)

Thanks for mentioning the PR, though there it refers to definitions in Zygote that seem to suggest Zygote should already know how to handle FFTs.

---

<div class="post-metadata">

**Author:** ![maedoc](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maedoc/32/8490_2.png) [@maedoc](https://discourse.julialang.org/u/maedoc)\
**Post date:** [October 22, 2021, 9:17pm UTC](https://discourse.julialang.org/t/gradients-of-fft/70223/5 "2021-10-22T21:17:43Z")

</div>

Trying out different algorithms listed in the docs, `sensalg=BacksolveAdjoint()` seems to have done the trick.

_edit_ this only works for the out of place case.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [October 23, 2021, 12:52pm UTC](https://discourse.julialang.org/t/gradients-of-fft/70223/6 "2021-10-23T12:52:30Z")

</div>

On the latest versions this all works out of the box.

> [@maedoc](#):
>
> _edit_ this only works for the out of place case.

Yes, because scalar tracing the FFT won’t work (at least without Enzyme support)