Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Real valued FFTs #1657

Merged
merged 6 commits into from
Jan 13, 2020
Merged

Real valued FFTs #1657

merged 6 commits into from
Jan 13, 2020

Conversation

shoyer
Copy link
Collaborator

@shoyer shoyer commented Nov 10, 2019

Note: The transpose rule is not correct yet (hence the failing tests).

Everything should be working -- this is ready for review!

Note: The transpose rule is not correct yet (hence the failing tests).
@joglekara
Copy link
Contributor

joglekara commented Dec 30, 2019

just want to pin this to the issue tracker #1877

@shoyer shoyer marked this pull request as ready for review January 12, 2020 00:35
@shoyer shoyer changed the title WIP: real valued fft functions rfft and irfft Jan 12, 2020
@shoyer shoyer requested a review from skye January 12, 2020 00:35
@shoyer shoyer changed the title rfft and irfft rfftn and irfftn Jan 12, 2020
Copy link
Collaborator

@mattjj mattjj left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Beautiful, thanks for doing this! I think it slipped off our radars.

jax/lax/lax_fft.py Show resolved Hide resolved
jax/lax/lax_fft.py Outdated Show resolved Hide resolved
@shoyer shoyer changed the title rfftn and irfftn Real valued FFTs Jan 13, 2020
@shoyer
Copy link
Collaborator Author

shoyer commented Jan 13, 2020

Now with rfft, irfft, rfft2 and irfft2...

self._CheckAgainstNumpy(onp_fn, np_fn, args_maker, check_dtypes=False,
tol=1e-4)
self._CompileAndCheck(np_fn, args_maker, check_dtypes=True)
# Test gradient for differentiable types.
if dtype in inexact_dtypes:
tol = 0.15 # TODO(skye): can we be more precise?
jtu.check_grads(np_fn, args_maker(), order=1, atol=tol, rtol=tol)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why get rid of order=1?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

order=2 already tests all derivatives up to order=2. So the separate order=1 was just redundant.

def _irfft_with_zeroed_inputs(irfft_fun):
# irfft isn't defined on the full domain of inputs, so in order to have a
# well defined derivative on the whole domain of the function, we zero-out
# the imaginary part of the first and possibly the last elements.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This logic is pretty complicated for a test... would be it be possible/simpler to limit the test inputs to irfft?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The challenge here is that in order for the finite difference calculation to be correct, I need to apply the restriction to the perturbation as well. I don't think our check_grads machinery is set-up to do that unless we modify the function passed into it.

Copy link
Member

@skye skye left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't give this a super close review, but overall LGTM. Lemme know if there's a certain part you wanted me to look at in particular.

@shoyer shoyer merged commit a5b6e8a into jax-ml:master Jan 13, 2020
@shoyer
Copy link
Collaborator Author

shoyer commented Jan 13, 2020

thanks @skye and @mattjj for reviews!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants