finufft: update to 2.4

.devcontainer/Dockerfile

@@ -0,0 +1 @@
+../ci/Dockerfile

.devcontainer/devcontainer.json

@@ -0,0 +1,9 @@
+{
+  "name": "CUDA",
+  "runArgs": [
+    "--gpus=all"
+  ],
+  "build": {
+    "dockerfile": "Dockerfile"
+  }
+}

CMakeLists.txt

@@ -41,9 +41,6 @@ if(JAX_FINUFFT_USE_CUDA)
 
         message(STATUS "jax_finufft: CUDA architectures: ${CMAKE_CUDA_ARCHITECTURES}")
 
-        # Propagate to finufft, because it doesn't look at CMAKE_CUDA_ARCHITECTURES by default
-        set(FINUFFT_CUDA_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES})
-
         # This needs to be run after the CMAKE_CUDA_ARCHITECTURES check, otherwise
         # it will set it to the compiler default
         enable_language(CUDA)

ci/Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu24.04
+FROM nvidia/cuda:12.9.0-cudnn-devel-ubuntu24.04
 
 RUN apt-get update && \
     DEBIAN_FRONTEND=noninteractive apt-get install -y \
@@ -8,4 +8,4 @@ RUN apt-get update && \
     curl \
     libfftw3-dev
 
-COPY --from=ghcr.io/astral-sh/uv:0.6.17 /uv /uvx /bin/
+COPY --from=ghcr.io/astral-sh/uv:0.7.11 /uv /uvx /bin/

ci/Jenkinsfile

@@ -21,7 +21,7 @@ pipeline {
             }
             steps {
                 sh '''
-                    uv run --extra test pytest -n 8 tests/
+                    uv run --extra test pytest -n 8
                 '''
             }
         }
@@ -34,7 +34,7 @@ pipeline {
             steps {
                 // TODO: add "-n 8", but GPU kernels don't seem to be thread-safe
                 sh '''
-                    uv run --extra test --extra cuda12 pytest tests/
+                    uv run --extra test --extra cuda12-local pytest
                 '''
             }
         }

lib/jax_finufft_cpu.cc

@@ -115,14 +115,13 @@ NB_MODULE(jax_finufft_cpu, m) {
 
   nb::class_<finufft_opts> opts(m, "FinufftOpts");
   opts.def("__init__",
-           [](finufft_opts *self, bool modeord, bool chkbnds, int debug, int spread_debug,
-              bool showwarn, int nthreads, int fftw, int spread_sort, bool spread_kerevalmeth,
-              bool spread_kerpad, double upsampfac, int spread_thread, int maxbatchsize,
-              int spread_nthr_atomic, int spread_max_sp_size) {
+           [](finufft_opts *self, bool modeord, int debug, int spread_debug, bool showwarn,
+              int nthreads, int fftw, int spread_sort, bool spread_kerevalmeth, bool spread_kerpad,
+              double upsampfac, int spread_thread, int maxbatchsize, int spread_nthr_atomic,
+              int spread_max_sp_size) {
              new (self) finufft_opts;
              default_opts<double>(self);
              self->modeord = int(modeord);
-             self->chkbnds = int(chkbnds);
              self->debug = debug;
              self->spread_debug = spread_debug;
              self->showwarn = int(showwarn);

pyproject.toml

@@ -19,6 +19,7 @@ dynamic = ["version"]
 [project.optional-dependencies]
 test = ["pytest", "pytest-xdist", "absl-py"]
 cuda12 = ["jax[cuda12]"]
+cuda12-local = ["jax[cuda12-local]"]
 
 [tool.scikit-build]
 metadata.version.provider = "scikit_build_core.metadata.setuptools_scm"

src/jax_finufft/options.py

@@ -44,7 +44,6 @@ class Opts:
     # These correspond to the default cufinufft options
     # set in vendor/finufft/src/cuda/cufinufft.cu
     modeord: bool = False
-    chkbnds: bool = True
     debug: DebugLevel = DebugLevel.Silent
     spread_debug: DebugLevel = DebugLevel.Silent
     showwarn: bool = False
@@ -77,7 +76,6 @@ def to_finufft_opts(self):
         compiled_with_omp = jax_finufft_cpu._omp_compile_check()
         return jax_finufft_cpu.FinufftOpts(
             self.modeord,
-            self.chkbnds,
             int(self.debug),
             int(self.spread_debug),
             self.showwarn,