Chain rules for FFT plans via AdjointPlans#67
Merged
devmotion merged 28 commits intoJuliaMath:masterfrom Jul 5, 2023
Merged
Conversation
Loading
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
An
rfftcan be written asPFwhereFis then x nFourier transform andPis a projection operator that removes the redundant information due to conjuagate symmetry. Because ofP, the adjoint of real FFTs (real inverse FFTs) require a special scaling before (after) applying the backwards transformation. As discussed in #63 this motivates supporting theBase.adjointoperation for plans to simplify the writing of backward rules for AD.The following functions must be implemented by backends in order for
output_size(p::Plan)andAdjointPlanto work:projection_style(p::Plan)which can either be:none,:real, or:real_inv.irfft_dim(p::Plan), only for those plans with:real_invprojection style, which gives the original length of the halved dimension.Using the adjoint plan, we can simplify the writing of backwards rules. I test the adjoint plans both directly and indirectly through tests of the rrule's.
NB: The interface has changed since the initial PR message. See the updated implementation docs in the PR for accurate info.