[Bug] Derivatives are wrong

[unalmis]

An algorithm I have written which has application in stellarator optimization would be improved by replacing some matrix multiplication transforms with type 2 non-uniform fast Fourier transforms. This algorithm is forward and reverse mode differentiable. Therefore, I would like to use jax-finufft as the flatiron institute has a nice implementation of this transform. However, my testing shows that the vector Jacobian products that jax-finufft assigns to these transformations are incorrect.

For example, here is a function $f$ which uses nufft2. Below I compute $f$ multiple times at different values of its argument $\lambda$.
Clearly the derivative at the marked point should match the derivative of the best linear approximation at that point. However, the derivative computed using the vector Jacobian product from jax-finufft is incorrect by four orders of magnitude.

This is an issue with jax-finufft; when I remove nuffts from the computation and use matrix multiplication transforms instead, vanilla jax computes the derivative correctly.

Illustration

Below is an illustration with the reverse mode derivative computation.

Reproduction

git clone https://github.com/PlasmaControl/DESC.git
cd DESC
git switch ku/nufft

conda create --name desc-env 'python>=3.10, <=3.13'
conda activate desc-env
pip install --editable .
pip install -r devtools/dev-requirements.txt

Now run the following command.

pytest -v --mpl -k test_binormal_drift_bounce2d

This command will run two tests which each check correctness of the computation and the derivative. The first test computes the algorithm using matrix multiplication, without any reference to jax-finufft. The second test uses nufft2 from jax-finufft with the options detailed here.

In this second test, if jax-finnuft computes the derivatives incorrectly, the test will pass. To recover the more natural behavior of the test failing when the derivatives are incorrect, one should move the code out of the with pytest.raises(AssertionError) block here.

[dfm]

Please put together a minimal, reproducible example that demonstrates the problem. This shouldn't depend on external projects, and it should be possible to directly share a code snippet here. We do have unit tests for gradients, but the coverage isn't wide so there could well be issues.

[unalmis]

Hi, thank you for the reply.

Reducing the function to something that simple is not possible. The algorithm to compute $f$ has more source code than the jax-finufft library. More demanding is if I were to rewrite a code snippet that does not depend on external projects I would need to rewrite most of jax's and DESC's source code.

Given that there are tests for some functions, it is probably unproductive to attempt guessing functions $g$ where jax-finufft fails to compute the derivative which are also simple enough for a MWE without external code. If I invest that time and discover such a function $g$, I will make a new issue. If it is believed the coverage is narrow enough that we might have luck with this guessing approach please let me know.

Given that I do not currently know how to invent a more minimal MWE for the function $f$, a different debugging strategy is required. The first purpose of a MWE is a sanity check for the developer to verify the source of the bug is not in an external library. This is covered in the issue description:

The computation of the nufft2s are correct. To test correctness of gradients, we can always compare to finite differencing the output of nufft2. In the above plot for example, it is clear the finite difference of $f$ computed by jax-finufft does not match the gradient that jax-finufft computes.

The second purpose of a MWE is to help isolate the issue to fix the debugging. It looks like the derivative assignment in jax-finufft is 50 lines of code, so that seems less intimidating to manipulate than inventing a MWE for $f$. If there is documentation that would assist in my debugging of this, please let me know. The first thing I feel drawn to check is whether the reshaping done for computing the function is consistent with the reshaping done in computing the derivative.

[andycasey]

Maybe I'm missing something here, but if you ran your own (non-minimal) reproducible example with some debugger that dropped you in to the stack where the test failed, wouldn't you simply be able to save the relevant local variables that are being supplied to nufft2 and jax-finufft to some memory-mapped arrays? Then you'd be able to create a ~10 line example that loads in those arrays and demonstrates that the gradients are not what you expect.

I started to run your (non-minimal) example out of interest because I was surprised if the gradients were wrong, but I stopped when I saw how big the DESC repository was.

[unalmis]

Hi Andy, thank you for the suggestion.

It is somewhat more complex than that. I had typed a few reasons why... maybe it suffices to mention that, if you cut the computation there, the derivative you see would be a large tensor. Not as fun for debugging as the scalar-valued derivative in the original problem.

Anyway, I had success with the guessing approach. I have added tests to the testing suite for simple functions whose derivatives are computed incorrectly by jax-finufft in #159 . We should debug those first.

[TomHilder]

Hi there, I think #160 should resolve this issue as well as the example in the PR #159. jvp was assuming the default mode ordering but I can see in your linked code you were using the alternate one. Thanks!

[unalmis]

The tests that verified my bugfix for this and #161 no longer pass in #159, so I cannot approve whether the issue was resolved. I thought I'd share since, if the bug was not fixed, it will not be caught on my end in a downstream repo because I implemented a patch to workaround the issue.