Remove FINUFFT warning

[unalmis]

Replaces jax-finufft's warning with the one desc developers want.

The base branch was changed.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    3.96 %    |     3.898e+03      |     4.052e+03      |    154.45    |       40.87        |       36.66        |
  test_proximal_jac_w7x_with_eq_update   |    2.69 %    |     6.557e+03      |     6.734e+03      |    176.61    |       163.22       |       161.89       |
  test_proximal_freeb_jac                |    0.84 %    |     1.315e+04      |     1.326e+04      |    110.46    |       83.64        |       83.74        |
  test_proximal_freeb_jac_blocked        |   -0.87 %    |     7.531e+03      |     7.466e+03      |    -65.18    |       74.31        |       73.91        |
  test_proximal_freeb_jac_batched        |    0.18 %    |     7.434e+03      |     7.447e+03      |    13.27     |       75.03        |       73.39        |
  test_proximal_jac_ripple               |   -0.94 %    |     3.476e+03      |     3.443e+03      |    -32.77    |       67.22        |       65.99        |
  test_proximal_jac_ripple_bounce1d      |    1.98 %    |     3.561e+03      |     3.631e+03      |    70.46     |       78.66        |       75.63        |
  test_eq_solve                          |    3.60 %    |     2.015e+03      |     2.088e+03      |    72.61     |       94.38        |       92.56        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.55%. Comparing base (52c31b7) to head (4690b05).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1944      +/-   ##
==========================================
+ Coverage   94.52%   94.55%   +0.02%     
==========================================
  Files         102      102              
  Lines       28785    28771      -14     
==========================================
- Hits        27210    27205       -5     
+ Misses       1575     1566       -9     

Files with missing lines	Coverage Δ
desc/integrals/_interp_utils.py	98.40% <100.00%> (+1.00%)	⬆️
desc/objectives/_fast_ion.py	100.00% <100.00%> (+3.19%)	⬆️
desc/objectives/_neoclassical.py	100.00% <100.00%> (+2.91%)	⬆️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[dpanici]

not approving yet, running into an error locally, will try to see the issue and post

checking further

[f0uriest]

The test folder isn't always included, ie if you install from pip. I don't think we need to run the full test anyways, we just need to test that it jits ok on gpu

[unalmis]

I prefer to do better than "if it compiles, ship it"

[dpanici]

Running

from desc.objectives import *

gives this error (on a clean install on cpu just doing pip install --editable ., I'm not sure if this is the way, in general, to safely grab things from the tests directory?

Traceback (most recent call last):
  File "/Users/dpanici/Research/work/desc_work/random/flexible_eq/try_import.py", line 1, in <module>
    from desc.objectives import *
  File "/Users/dpanici/Research/DESC/desc/objectives/__init__.py", line 3, in <module>
    from ._bootstrap import BootstrapRedlConsistency
  File "/Users/dpanici/Research/DESC/desc/objectives/_bootstrap.py", line 6, in <module>
    from desc.compute import get_profiles, get_transforms
  File "/Users/dpanici/Research/DESC/desc/compute/__init__.py", line 30, in <module>
    from . import (
    ...<15 lines>...
    )
  File "/Users/dpanici/Research/DESC/desc/compute/_bootstrap.py", line 16, in <module>
    from ..integrals.surface_integral import surface_averages_map
  File "/Users/dpanici/Research/DESC/desc/integrals/__init__.py", line 3, in <module>
    from ._bounce_utils import fast_chebyshev, fast_cubic_spline, fourier_chebyshev
  File "/Users/dpanici/Research/DESC/desc/integrals/_bounce_utils.py", line 8, in <module>
    from desc.integrals._interp_utils import (
    ...<8 lines>...
    )
  File "/Users/dpanici/Research/DESC/desc/integrals/_interp_utils.py", line 998, in <module>
    _test_gpu_jax_finufft()
    ~~~~~~~~~~~~~~~~~~~~~^^
  File "/Users/dpanici/Research/DESC/desc/integrals/_interp_utils.py", line 979, in _test_gpu_jax_finufft
    from tests.test_interp_utils import TestFastInterp, _test_inputs_1D
ModuleNotFoundError: No module named 'tests'

[YigitElma]

I still think these kind of checks that include computation must be application specific. For example, on my laptop, I use GPU but the installation instruction doesn't work, so I just installed jax-finufft cpu to shut the warning down, but with this, I will always get this warning. The additional import time is also not necessary.

Also, we can either use this PR or some other one, but the installation instructions are not complete for jax-finufft.

• For example, the local GPU one doesn't work. You need to install CudaToolkit, we should at least mention that (I haven't tested it with Cuda Toolkit because it might cause my other environments to fail with some CUDA related errors, I don't want to take that risk).
• Della instructions only work on della-gpu login node, this again needs explanation.
• Another thing is that if you try to use the environment with a H100, you cannot, you need to compile it again in a different environment with CMAKE_CUDA_ARCHITECTURES=90 or something like that. Because on Della and Perlmutter, login nodes have A100 with 8.0 architecture.
• I shared the Perlmutter instructions in #1937. They are very very limiting but work. I tried making it more general, but looks like if you change one line of the instructions, everything fails.

[dpanici]

If in our current install instructions, running this is enough to confirm what this warning is checking explicitly, then I think this PR may not be necessary.

If it is not enough, then just extending the docs to include the extra check on grad of nufft1d2r or whatever should be sufficient. At least in the opinions of myself, @ddudt and @YigitElma

This only really needs to be checked once upon install, not everytime the code is ran. @f0uriest if you think differently let us know since I think you had originally pointed out the desire for this warning?

from desc import set_device
set_device("gpu")

from desc.examples import get
from desc.objectives import ObjectiveFunction, GammaC

obj = ObjectiveFunction(GammaC(get("W7-X"), num_transit=1, num_pitch=1))
obj.build()
x = obj.x()
obj.compute_scaled_error(x).block_until_ready()

[dpanici]

Ideally actually, can this be done within the objectives/compute functions which use jax-finufft? instead of being upon import or in installation?

I have made the changes discussed in the dev meeting. Other people should review

[ddudt]

I like this idea of making the user test their installation instead of checking it every time the functions are imported. But we should add a note in the Verifying your Installation section with example code of how to test it. Otherwise this looks good.

[YigitElma]

I like this idea of making the user test their installation instead of checking it every time the functions are imported. But we should add a note in the Verifying your Installation section with example code of how to test it. Otherwise this looks good.

I think we already have that

from desc import set_device
set_device("gpu")

from desc.examples import get
from desc.objectives import ObjectiveFunction, GammaC

obj = ObjectiveFunction(GammaC(get("W7-X"), num_transit=1, num_pitch=1))
obj.build()
x = obj.x()
obj.compute_scaled_error(x).block_until_ready()

This is under the CPU+GPU tab of verify your installation. CPU version should work out of the box.

[dpanici]

flatironinstitute/jax-finufft#206
FYI