Incorrect pullback of `AbstractFFTs.[i]rfft`?

[trahflow]

Zygote defines an adjoint of AbstractFFTs.[i]rfft here.
On the other hand, AbstractFFTs defines rrules for [i]rfft and other functions itself (implemented by @devmotion ) here

I assume that Zygote prefers its own adjoint over the ChainRulesCore.rrule of AbstractFFTs (at least that is what I think I'm seeing).

Now the problem is that the two implementations differ (not just by the additional projection into the real domain, but by additional scaling factors) and give different results.
I haven't yet taken the time to figure out how the scaling factors in AbstractFFTs are motivated (unfortunately there's no comment that gives a hint), but finite differencing seems to agree with AbstractFFTs implementation and disagree with Zygote's.

Side question: Can I somehow tell Zygote to prefer the ChainRulesCore.rrule over its own @adjoint?

[ToucheSir]

#1158 I believe is the same issue by another name. #899 looks related. I think #1173 was meant to be a fix, but have no expertise to evaluate it.

Side question: Can I somehow tell Zygote to prefer the ChainRulesCore.rrule over its own @adjoint?

Not without resorting to piracy, no. I think at this point we should just remove all FFT-related rules from Zygote and let the appropriate chainrules take over, but that would require someone with the requisite experience on how to do so.

[CarloLucibello]

I think we should remove Zygote's adjoints for all ffts rule. If we also manage to remove the pirate methods for Fill compatibility we could ditch the AbstractFFTs dependency entirely. @trahflow would you mind filing a PR?

[trahflow]

#1158 I believe is the same issue by another name

Sorry for the duplicate. I did a search but probably only for "irfft".

@trahflow would you mind filing a PR?

I will give it a shot. I won't get around doing it before the end of the week though.

[trahflow]

AbstractFFTs does currently not yet define rules for fft plans. There's this PR: JuliaMath/AbstractFFTs.jl#67 which provides them. One would probably want to get this merged first?

[ToucheSir]

Makes sense, I guess that's where the last chain of PRs got stuck?

[gaurav-arya]

The plan PR should get merged in the coming weeks when I get time to address recent review comments. But I don't see why that's an obstacle to getting rid of Zygote fft/rfft/etc. rules asap if they still exist, and just leaving the (also incorrect in general) plan rules around until they are replaced upstream?

[ToucheSir]

I have no idea what that entails, but it seems easy to verify existing tests pass if anyone wants to throw up a quick PR.

[trahflow]

But I don't see why that's an obstacle to getting rid of Zygote fft/rfft/etc. rules asap if they still exist, and just leaving the (also incorrect in general) plan rules around until they are replaced upstream?

I gave it a shot here:
#1386

[trahflow]

But I don't see why that's an obstacle to getting rid of Zygote fft/rfft/etc. rules asap if they still exist, and just leaving the (also incorrect in general) plan rules around until they are replaced upstream?

@gaurav-arya Note that if the dims argument is not provided to fft() etc., the rules provided by JuliaMath/AbstractFFTs.jl#58 aren't actually hit, and the compilation will fail trying to differentiate a try/except block in AbstractFFTs.plan_fft() once those rules are removed from Zygote.

This is certainly better than silently giving incorrect results as in [i]rfft(), but probably worse for fft() etc. where the rules in Zygote were seemingly correct.

[trahflow]

Fixed by #1386 and JuliaMath/AbstractFFTs.jl#84

The adjoints involving AbstractFFTs.Plan are still incorrect but that should be tracked by another issue I think (maybe extend #899 ?)