Repository: flatironinstitute/jax-finufft (language: python)

Issue:
MLIR translation rule for primitive 'nufft2' not found for platform cuda
Thank you for your work on jax-finufft. I am successfully using it targeting my cpu, but am having issues with it when using my gpu. Bascially, I get a "NotImplementedError: MLIR translation rule for primitive 'nufft2' not found for platform cuda" as soon as jax tries to do autodifferentiation.

As an example:
`
$ pytest ops_test.py::test_nufft1_forward[2-False-50-75--1]
========================== test session starts ===========================
platform linux -- Python 3.10.12, pytest-7.4.2, pluggy-1.3.0
rootdir: <PATH>/jax-finufft/tests
plugins: anyio-3.6.2
collected 1 item                                                         

ops_test.py F                                                      [100%]

================================ FAILURES ================================
_________________ test_nufft1_forward[2-False-50-75--1] __________________

    @pytest.mark.parametrize(
        "ndim, x64, num_nonnuniform, num_uniform, iflag",
        product([1, 2, 3], [False, True], [50], [75], [-1, 1]),
    )
    def test_nufft1_forward(ndim, x64, num_nonnuniform, num_uniform, iflag):
        if ndim == 1 and jax.default_backend() != "cpu":
            pytest.skip("1D transforms not implemented on GPU")
    
        random = np.random.default_rng(657)
    
        eps = 1e-10 if x64 else 1e-7
        dtype = np.double if x64 else np.single
        cdtype = np.cdouble if x64 else np.csingle
    
        num_uniform = tuple(num_uniform // ndim + 5 * np.arange(ndim))
        ks = [np.arange(-np.floor(n / 2), np.floor((n - 1) / 2 + 1)) for n in num_uniform]
    
        x = random.uniform(-np.pi, np.pi, size=(ndim, num_nonnuniform)).astype(dtype)
        c = random.normal(size=num_nonnuniform) + 1j * random.normal(size=num_nonnuniform)
        c = c.astype(cdtype)
        f_expect = np.zeros(num_uniform, dtype=cdtype)
        for coords in product(*map(range, num_uniform)):
            k_vec = np.array([k[n] for (n, k) in zip(coords, ks)])
            f_expect[coords] = np.sum(c * np.exp(iflag * 1j * np.dot(k_vec, x)))
    
        with jax.experimental.enable_x64(x64):
>           f_calc = nufft1(num_uniform, c, *x, eps=eps, iflag=iflag)

ops_test.py:45: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

    @partial(jit, static_argnums=(0,), static_argnames=("iflag", "eps"))
    def nufft1(output_shape, source, *points, iflag=1, eps=1e-6):
        iflag = int(iflag)
        eps = float(eps)
        ndim = len(points)
        if not 1 <= ndim <= 3:
            raise ValueError("Only 1-, 2-, and 3-dimensions are supported")
    
        # Support passing a scalar output_shape
        output_shape = np.atleast_1d(output_shape).astype(np.int64)
        if output_shape.shape != (ndim,):
            raise ValueError(f"output_shape must have shape: ({ndim},)")
    
        # Handle broadcasting and reshaping of inputs
        index, source, *points = shapes.broadcast_and_flatten_inputs(
            output_shape, source, *points
        )
    
        # Execute the transform primitive
>       result = nufft1_p.bind(
            source, *points, output_shape=output_shape, iflag=iflag, eps=eps
        )
E       jax._src.source_info_util.JaxStackTraceBeforeTransformation: NotImplementedError: MLIR translation rule for primitive 'nufft1' not found for platform cuda
E       
E       The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
E       
E       --------------------

../../../../.local/lib/python3.10/site-packages/jax_finufft/ops.py:33: JaxStackTraceBeforeTransformation

The above exception was the direct cause of the following exception:

ndim = 2, x64 = False, num_nonnuniform = 50, num_uniform = (37, 42)
iflag = -1

    @pytest.mark.parametrize(
        "ndim, x64, num_nonnuniform, num_uniform, iflag",
        product([1, 2, 3], [False, True], [50], [75], [-1, 1]),
    )
    def test_nufft1_forward(ndim, x64, num_nonnuniform, num_uniform, iflag):
        if ndim == 1 and jax.default_backend() != "cpu":
            pytest.skip("1D transforms not implemented on GPU")
    
        random = np.random.default_rng(657)
    
        eps = 1e-10 if x64 else 1e-7
        dtype = np.double if x64 else np.single
        cdtype = np.cdouble if x64 else np.csingle
    
        num_uniform = tuple(num_uniform // ndim + 5 * np.arange(ndim))
        ks = [np.arange(-np.floor(n / 2), np.floor((n - 1) / 2 + 1)) for n in num_uniform]
    
        x = random.uniform(-np.pi, np.pi, size=(ndim, num_nonnuniform)).astype(dtype)
        c = random.normal(size=num_nonnuniform) + 1j * random.normal(size=num_nonnuniform)
        c = c.astype(cdtype)
        f_expect = np.zeros(num_uniform, dtype=cdtype)
        for coords in product(*map(range, num_uniform)):
            k_vec = np.array([k[n] for (n, k) in zip(coords, ks)])
            f_expect[coords] = np.sum(c * np.exp(iflag * 1j * np.dot(k_vec, x)))
    
        with jax.experimental.enable_x64(x64):
>           f_calc = nufft1(num_uniform, c, *x, eps=eps, iflag=iflag)

ops_test.py:45: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../../../.local/lib/python3.10/site-packages/jax/_src/pjit.py:238: in cache_miss
    outs, out_flat, out_tree, args_flat = _python_pjit_helper(
../../../../.local/lib/python3.10/site-packages/jax/_src/pjit.py:185: in _python_pjit_helper
    out_flat = pjit_p.bind(*args_flat, **params)
../../../../.local/lib/python3.10/site-packages/jax/_src/core.py:2592: in bind
    return self.bind_with_trace(top_trace, args, params)
../../../../.local/lib/python3.10/site-packages/jax/_src/core.py:363: in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
../../../../.local/lib/python3.10/site-packages/jax/_src/core.py:817: in process_primitive
    return primitive.impl(*tracers, **params)
../../../../.local/lib/python3.10/site-packages/jax/_src/pjit.py:1229: in _pjit_call_impl
    compiled = _pjit_lower(
../../../../.local/lib/python3.10/site-packages/jax/_src/pjit.py:1315: in _pjit_lower
    return _pjit_lower_cached(jaxpr, in_shardings, out_shardings, *args, **kwargs)
../../../../.local/lib/python3.10/site-packages/jax/_src/pjit.py:1374: in _pjit_lower_cached
    return pxla.lower_sharding_computation(
../../../../.local/lib/python3.10/site-packages/jax/_src/profiler.py:314: in wrapper
    return func(*args, **kwargs)
../../../../.local/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py:2551: in lower_sharding_computation
    lowering_result = mlir.lower_jaxpr_to_module(
../../../../.local/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:742: in lower_jaxpr_to_module
    lower_jaxpr_to_fun(
../../../../.local/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:1037: in lower_jaxpr_to_fun
    out_vals, tokens_out = jaxpr_subcomp(ctx.replace(name_stack=callee_name_stack),
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

ctx = ModuleContext(context=<jaxlib.mlir._mlir_libs._site_initialize.<locals>.Context object at 0x7f662850d170>, module=<jax...annel_iterator=count(1), host_callbacks=[], dim_vars=(), cached_primitive_lowerings={}, cached_call_jaxpr_lowerings={})
jaxpr = { lambda ; a:c64[50] b:f32[50] c:f32[50]. let
    d:c64[1,1,50] = reshape[dimensions=None new_sizes=(1, 1, 50)] a
    ...1e-07 iflag=-1 output_shape=[37 42]] d e f
    h:c64[37,42] = reshape[dimensions=None new_sizes=(37, 42)] g
  in (h,) }
tokens = <jax._src.interpreters.mlir.TokenSet object at 0x7f66284e6c50>
consts = [], dim_var_values = []
args = ([<jaxlib.mlir._mlir_libs._mlir.ir.BlockArgument object at 0x7f66285118f0>], [<jaxlib.mlir._mlir_libs._mlir.ir.BlockArgument object at 0x7f6628511930>], [<jaxlib.mlir._mlir_libs._mlir.ir.BlockArgument object at 0x7f66285119b0>])
aval = <function jaxpr_subcomp.<locals>.aval at 0x7f66284db880>
write = <function jaxpr_subcomp.<locals>.write at 0x7f66284db9a0>

    def jaxpr_subcomp(ctx: ModuleContext, jaxpr: core.Jaxpr,
                      tokens: TokenSet,
                      consts: Sequence[Sequence[ir.Value]],
                      *args: Sequence[ir.Value],
                      dim_var_values: Sequence[ir.Value]
                      ) -> Tuple[Sequence[Sequence[ir.Value]], TokenSet]:
      """Lowers a jaxpr into MLIR, inlined into an existing function.
    
      Assumes that an MLIR context, location, and insertion point are set.
    
      dim_var_values: the list of dimension variables values in the current
        IR function, in the order of ctx.dim_vars.
      """
      assert ctx.platform != "gpu"
      def read(v: core.Atom) -> Sequence[ir.Value]:
        if type(v) is core.Literal:
          return ir_constants(v.val, canonicalize_types=True)
        else:
          assert isinstance(v, core.Var)
          return env[v]
    
      def aval(v: core.Atom) -> core.AbstractValue:
        if type(v) is core.Literal:
          return xla.abstractify(v.val)
        else:
          return v.aval
    
      def write(v: core.Var, node: Sequence[ir.Value]):
        assert node is not None
        env[v] = tuple(node)
    
    
      env: Dict[core.Var, Tuple[ir.Value, ...]] = {}
    
      assert len(args) == len(jaxpr.invars), (jaxpr, args)
      assert len(consts) == len(jaxpr.constvars), (jaxpr, consts)
      assert all(isinstance(v, ir.Value) for vs in consts for v in vs), consts
      assert len(ctx.dim_vars) == len(dim_var_values), (ctx.dim_vars, dim_var_values)
      map(write, jaxpr.constvars, consts)
      map(write, jaxpr.invars, args)
      for eqn in jaxpr.eqns:
        in_nodes = map(read, eqn.invars)
        assert isinstance(ctx.name_stack, source_info_util.NameStack), type(ctx.name_stack)
        source_info = eqn.source_info.replace(
            name_stack=ctx.name_stack + eqn.source_info.name_stack)
        loc = _source_info_to_location(eqn.primitive, eqn.params, source_info,
                                       ctx.name_stack)
        with source_info_util.user_context(eqn.source_info.traceback), loc:
          if eqn.primitive in _platform_specific_lowerings[ctx.platform]:
            rule = _platform_specific_lowerings[ctx.platform][eqn.primitive]
          elif eqn.primitive in xla._backend_specific_translations[ctx.platform]:
            rule = xla_fallback_lowering(eqn.primitive)
          elif eqn.primitive in _lowerings:
            rule = _lowerings[eqn.primitive]
          elif eqn.primitive in xla._translations:
            rule = xla_fallback_lowering(eqn.primitive)
          else:
>           raise NotImplementedError(
                f"MLIR translation rule for primitive '{eqn.primitive.name}' not "
                f"found for platform {ctx.platform}")
E           jax._src.traceback_util.UnfilteredStackTrace: NotImplementedError: MLIR translation rule for primitive 'nufft1' not found for platform cuda
E           
E           The stack trace below excludes JAX-internal frames.
E           The preceding is the original exception that occurred, unmodified.
E           
E           --------------------

../../../../.local/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:1155: UnfilteredStackTrace

The above exception was the direct cause of the following exception:

ndim = 2, x64 = False, num_nonnuniform = 50, num_uniform = (37, 42)
iflag = -1

    @pytest.mark.parametrize(
        "ndim, x64, num_nonnuniform, num_uniform, iflag",
        product([1, 2, 3], [False, True], [50], [75], [-1, 1]),
    )
    def test_nufft1_forward(ndim, x64, num_nonnuniform, num_uniform, iflag):
        if ndim == 1 and jax.default_backend() != "cpu":
            pytest.skip("1D transforms not implemented on GPU")
    
        random = np.random.default_rng(657)
    
        eps = 1e-10 if x64 else 1e-7
        dtype = np.double if x64 else np.single
        cdtype = np.cdouble if x64 else np.csingle
    
        num_uniform = tuple(num_uniform // ndim + 5 * np.arange(ndim))
        ks = [np.arange(-np.floor(n / 2), np.floor((n - 1) / 2 + 1)) for n in num_uniform]
    
        x = random.uniform(-np.pi, np.pi, size=(ndim, num_nonnuniform)).astype(dtype)
        c = random.normal(size=num_nonnuniform) + 1j * random.normal(size=num_nonnuniform)
        c = c.astype(cdtype)
        f_expect = np.zeros(num_uniform, dtype=cdtype)
        for coords in product(*map(range, num_uniform)):
            k_vec = np.array([k[n] for (n, k) in zip(coords, ks)])
            f_expect[coords] = np.sum(c * np.exp(iflag * 1j * np.dot(k_vec, x)))
    
        with jax.experimental.enable_x64(x64):
>           f_calc = nufft1(num_uniform, c, *x, eps=eps, iflag=iflag)
E           NotImplementedError: MLIR translation rule for primitive 'nufft1' not found for platform cuda

ops_test.py:45: NotImplementedError
======================== short test summary info =========================
FAILED ops_test.py::test_nufft1_forward[2-False-50-75--1] - NotImplementedError: MLIR translation rule for primitive 'nufft1' not...
=========================== 1 failed in 0.86s ============================

`


Discussion / hints:
For context, I am on a Thinkpad P1 Gen 2 running Ubuntu 22.04 with CUDA 12.3.
`
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 545.29.06              Driver Version: 545.29.06    CUDA Version: 12.3     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Quadro T2000                   Off | 00000000:01:00.0 Off |                  N/A |
| N/A   44C    P8               4W /  30W |   3060MiB /  4096MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      1499      G   /usr/lib/xorg/Xorg                            4MiB |
|    0   N/A  N/A     21281      C   /usr/bin/python3                           3052MiB |
+---------------------------------------------------------------------------------------+
`


I have installed finufft (from source, built following [these] instructions, which were incomplete and did not include installation and path updates)

1. `mkdir build`
2. `cd build`
3. `cmake -D FINUFFT_USE_CUDA=ON -D CMAKE_CUDA_ARCHITECTURES=75 ..`
4. `cmake --build .`
5. `cmake --install . # this was missing from finufft instructions`
6. added `/usr/local/lib` to `LD_LIBRARY_PATH`, also missing from finufft instructions
7. `pip install python/cufinufft`
8. `pytest --framework=pycuda python/cufinufft/tests
========================== test session starts ===========================
platform linux -- Python 3.10.12, pytest-7.4.2, pluggy-1.3.0
plugins: anyio-3.6.2
collected 666 items                                                      

python/cufinufft/tests/test_array_ordering.py .                    [  0%]
python/cufinufft/tests/test_basic.py ............................. [  4%]
.................................................................. [ 14%]
.................................................................. [ 24%]
........................................................           [ 32%]
python/cufinufft/tests/test_error_checks.py ........               [ 33%]
python/cufinufft/tests/test_examples.py .ss.s.s                    [ 34%]
python/cufinufft/tests/test_multi.py s                             [ 35%]
python/cufinufft/tests/test_simple.py ............................ [ 39%]
.................................................................. [ 49%]
.................................................................. [ 59%]
.................................................................. [ 69%]
.................................................................. [ 78%]
.................................................................. [ 88%]
.................................................................. [ 98%]
........                                                           [100%]

============================ warnings summary ============================
tests/test_array_ordering.py::test_type1_ordering[pycuda]
  /home/<home>/.local/lib/python3.10/site-packages/pycuda/compyte/dtypes.py:120: DeprecationWarning: `np.bool8` is a deprecated alias for `np.bool_`.  (Deprecated NumPy 1.24)
    reg.get_or_register_dtype("bool", np.bool8)

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=============== 661 passed, 5 skipped, 1 warning in 3.90s ================

An error like `NotImplementedError: MLIR translation rule for primitive 'nufft1' not found for platform cuda` usually means that either jax-finufft wasn't compiled with CUDA support, or that it was compiled but the extension can't be loaded (e.g. #46).

Can you share the exact commands you're running to build jax-finufft and their output (and make the pip install verbose with `python -m pip install -v .`)? Note that there's a section in the instructions specific to building with CUDA support. I would also double-check that `LD_LIBRARY_PATH` is set per the instructions.

finufft/cufinufft are bundled with jax-finufft, so you don't need to install them separately. I don't think jax-finufft has any way to see them, though, so it's probably harmless.
Thanks for the reply, and sorry it's taken me this long to circle back.

Indeed, there was an environment variable passing problem that stemmed from some unclarities in the install documentation. Once we get this working, I can offer a documentation of my process.

Build and installation is resolved, but now get a straight segfault on when running tests. I'm still trying to track down where that segfault comes from, and any help would be appreciated. Here is where I am now:

$ cd jax-finufft
$ nvidia-smi --query-gpu=compute_cap --format=csv,noheader
7.5
$ export CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=75 -DJAX_FINUFFT_USE_CUDA=ON"
$ echo $LD_LIBRARY_PATH
/usr/local/cuda-12.3/lib64:/usr/local/lib:
$ rm -rf build
$ pip install pybind11[global]
$ export CMAKE_PREFIX_PATH=/home/aparsons/.local/lib/python3.10/site-packages/pybind11/share/cmake/pybind11
$ pip install -v .
>
Using pip 24.0 from /home/aparsons/.local/lib/python3.10/site-packages/pip (python 3.10)
Defaulting to user installation because normal site-packages is not writeable
Processing /home/aparsons/projects/hera/jax-finufft
  Running command pip subprocess to install build dependencies
  Collecting pybind11>=2.6
    Using cached pybind11-2.11.1-py3-none-any.whl.metadata (9.5 kB)
  Collecting scikit-build-core>=0.5
    Using cached scikit_build_core-0.8.2-py3-none-any.whl.metadata (19 kB)
  Collecting exceptiongroup (from scikit-build-core>=0.5)
    Using cached exceptiongroup-1.2.0-py3-none-any.whl.metadata (6.6 kB)
  Collecting packaging>=20.9 (from scikit-build-core>=0.5)
    Using cached packaging-24.0-py3-none-any.whl.metadata (3.2 kB)
  Collecting tomli>=1.1 (from scikit-build-core>=0.5)
    Using cached tomli-2.0.1-py3-none-any.whl.metadata (8.9 kB)
  Using cached pybind11-2.11.1-py3-none-any.whl (227 kB)
  Using cached scikit_build_core-0.8.2-py3-none-any.whl (140 kB)
  Using cached packaging-24.0-py3-none-any.whl (53 kB)
  Using cached tomli-2.0.1-py3-none-any.whl (12 kB)
  Using cached exceptiongroup-1.2.0-py3-none-any.whl (16 kB)
  Installing collected packages: tomli, pybind11, packaging, exceptiongroup, scikit-build-core
  Successfully installed exceptiongroup-1.2.0 packaging-24.0 pybind11-2.11.1 scikit-build-core-0.8.2 tomli-2.0.1
  Installing build dependencies ... done
  Running command Getting requirements to build wheel
  Getting requirements to build wheel ... done
  Running command pip subprocess to install backend dependencies
  Collecting pathspec
    Using cached pathspec-0.12.1-py3-none-any.whl.metadata (21 kB)
  Collecting setuptools-scm
    Using cached setuptools_scm-8.0.4-py3-none-any.whl.metadata (6.4 kB)
  Collecting pyproject_metadata
    Using cached pyproject_metadata-0.7.1-py3-none-any.whl.metadata (3.0 kB)
  Collecting ninja>=1.5
    Using cached ninja-1.11.1.1-py2.py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl.metadata (5.3 kB)
  Collecting packaging>=20 (from setuptools-scm)
    Using cached packaging-24.0-py3-none-any.whl.metadata (3.2 kB)
  Collecting setuptools (from setuptools-scm)
    Using cached setuptools-69.2.0-py3-none-any.whl.metadata (6.3 kB)
  Collecting typing-extensions (from setuptools-scm)
    Using cached typing_extensions-4.10.0-py3-none-any.whl.metadata (3.0 kB)
  Collecting tomli>=1 (from setuptools-scm)
    Using cached tomli-2.0.1-py3-none-any.whl.metadata (8.9 kB)
  Using cached pathspec-0.12.1-py3-none-any.whl (31 kB)
  Using cached setuptools_scm-8.0.4-py3-none-any.whl (42 kB)
  Using cached pyproject_metadata-0.7.1-py3-none-any.whl (7.4 kB)
  Using cached ninja-1.11.1.1-py2.py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl (307 kB)
  Using cached packaging-24.0-py3-none-any.whl (53 kB)
  Using cached tomli-2.0.1-py3-none-any.whl (12 kB)
  Using cached setuptools-69.2.0-py3-none-any.whl (821 kB)
  Using cached typing_extensions-4.10.0-py3-none-any.whl (33 kB)
  Installing collected packages: ninja, typing-extensions, tomli, setuptools, pathspec, packaging, setuptools-scm, pyproject_metadata
  Successfully installed ninja-1.11.1.1 packaging-24.0 pathspec-0.12.1 pyproject_metadata-0.7.1 setuptools-69.2.0 setuptools-scm-8.0.4 tomli-2.0.1 typing-extensions-4.10.0
  Installing backend dependencies ... done
  Running command Preparing metadata (pyproject.toml)
  *** scikit-build-core 0.8.2 using CMake 3.22.1 (metadata_wheel)
  Preparing metadata (pyproject.toml) ... done
Requirement already satisfied: jax in /home/aparsons/.local/lib/python3.10/site-packages (from jax-finufft==0.0.4.dev53+g049f0b9) (0.4.8)
Requirement already satisfied: jaxlib in /home/aparsons/.local/lib/python3.10/site-packages (from jax-finufft==0.0.4.dev53+g049f0b9) (0.4.7+cuda12.cudnn88)
Requirement already satisfied: ml-dtypes>=0.0.3 in /home/aparsons/.local/lib/python3.10/site-packages (from jax->jax-finufft==0.0.4.dev53+g049f0b9) (0.0.4)
Requirement already satisfied: numpy>=1.21 in /home/aparsons/.local/lib/python3.10/site-packages (from jax->jax-finufft==0.0.4.dev53+g049f0b9) (1.25.0)
Requirement already satisfied: opt-einsum in /home/aparsons/.local/lib/python3.10/site-packages (from jax->jax-finufft==0.0.4.dev53+g049f0b9) (3.3.0)
Requirement already satisfied: scipy>=1.7 in /home/aparsons/.local/lib/python3.10/site-packages (from jax->jax-finufft==0.0.4.dev53+g049f0b9) (1.10.1)
Building wheels for collected packages: jax-finufft
  Running command Building wheel for jax-finufft (pyproject.toml)
  *** scikit-build-core 0.8.2 using CMake 3.22.1 (wheel)
  *** Configuring CMake...
  loading initial cache file build/cp310-cp310-manylinux_2_35_x86_64/CMakeInit.txt
  -- The C compiler identification is GNU 11.4.0
  -- The CXX compiler identification is GNU 11.4.0
  -- Detecting C compiler ABI info
  -- Detecting C compiler ABI info - done
  -- Check for working C compiler: /usr/bin/cc - skipped
  -- Detecting C compile features
  -- Detecting C compile features - done
  -- Detecting CXX compiler ABI info
  -- Detecting CXX compiler ABI info - done
  -- Check for working CXX compiler: /usr/bin/c++ - skipped
  -- Detecting CXX compile features
  -- Detecting CXX compile features - done
  -- Using CMake version: 3.22.1
  -- Looking for a CUDA compiler
  -- Looking for a CUDA compiler - /usr/local/cuda-12.3/bin/nvcc
  -- CUDA compiler found; compiling with GPU support
  -- The CUDA compiler identification is NVIDIA 12.3.107
  -- Detecting CUDA compiler ABI info
  -- Detecting CUDA compiler ABI info - done
  -- Check for working CUDA compiler: /usr/local/cuda-12.3/bin/nvcc - skipped
  -- Detecting CUDA compile features
  -- Detecting CUDA compile features - done
  -- Downloading CPM.cmake to /home/aparsons/projects/hera/jax-finufft/build/cp310-cp310-manylinux_2_35_x86_64/cmake/CPM_0.38.0.cmake
  -- CPM: Adding package findfftw@ (master)
  -- Found PkgConfig: /usr/bin/pkg-config (found version "0.29.2")
  -- Found FFTW: /usr/include
  -- Found OpenMP_C: -fopenmp (found version "4.5")
  -- Found OpenMP_CXX: -fopenmp (found version "4.5")
  -- Found OpenMP: TRUE (found version "4.5")
  -- Found CUDAToolkit: /usr/local/cuda-12.3/include (found version "12.3.107")
  -- Looking for pthread.h
  -- Looking for pthread.h - found
  -- Performing Test CMAKE_HAVE_LIBC_PTHREAD
  -- Performing Test CMAKE_HAVE_LIBC_PTHREAD - Success
  -- Found Threads: TRUE
  -- Found PythonInterp: /usr/bin/python3 (found suitable version "3.10.12", minimum required is "3.6")
  -- Found PythonLibs: /usr/lib/x86_64-linux-gnu/libpython3.10.so
  -- Found pybind11: /tmp/pip-build-env-9jy7iupp/overlay/local/lib/python3.10/dist-packages/pybind11/include (found version "2.11.1")
  -- Configuring done
  -- Generating done
  -- Build files have been written to: /home/aparsons/projects/hera/jax-finufft/build/cp310-cp310-manylinux_2_35_x86_64
  *** Building project with Ninja...
  [1/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f32.dir/fortran/finufftfort.cpp.o
  [2/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f64.dir/fortran/finufftfort.cpp.o
  [3/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f64.dir/src/utils.cpp.o
  [4/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f64.dir/src/simpleinterfaces.cpp.o
  [5/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f32.dir/src/utils.cpp.o
  [6/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f32.dir/src/simpleinterfaces.cpp.o
  [7/42] Building CXX object vendor/finufft/CMakeFiles/finufft.dir/src/utils_precindep.cpp.o
  [8/42] Building CXX object vendor/finufft/CMakeFiles/finufft_static.dir/contrib/legendre_rule_fast.cpp.o
  [9/42] Building CXX object vendor/finufft/CMakeFiles/finufft.dir/contrib/legendre_rule_fast.cpp.o
  [10/42] Building CXX object vendor/finufft/CMakeFiles/finufft_static.dir/src/utils_precindep.cpp.o
  [11/42] Building CXX object vendor/finufft/src/cuda/CMakeFiles/cufinufft_common_objects.dir/utils.cpp.o
  [12/42] Building CXX object vendor/finufft/src/cuda/CMakeFiles/cufinufft_common_objects.dir/__/__/contrib/legendre_rule_fast.cpp.o
  [13/42] Building CXX object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/spreadinterp.cpp.o
  [14/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f32.dir/src/finufft.cpp.o
  [15/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f64.dir/src/finufft.cpp.o
  [16/42] Building CUDA object CMakeFiles/jax_finufft_gpu.dir/lib/kernels.cc.cu.o
  [17/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/1d/cufinufft1d.cu.o
  [18/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f64.dir/src/spreadinterp.cpp.o
  [19/42] Building CXX object vendor/finufft/CMakeFiles/finufft_f32.dir/src/spreadinterp.cpp.o
  [20/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/2d/cufinufft2d.cu.o
  [21/42] Linking CXX static library vendor/finufft/libfinufft_static.a
  [22/42] Linking CXX shared library vendor/finufft/libfinufft.so
  [23/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/3d/cufinufft3d.cu.o
  [24/42] Building CXX object CMakeFiles/jax_finufft_gpu.dir/lib/jax_finufft_gpu.cc.o
  [25/42] Building CXX object CMakeFiles/jax_finufft_cpu.dir/lib/jax_finufft_cpu.cc.o
  [26/42] Linking CXX shared module jax_finufft_cpu.cpython-310-x86_64-linux-gnu.so
  [27/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/memtransfer_wrapper.cu.o
  [28/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/cufinufft.cu.o
  [29/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/deconvolve_wrapper.cu.o
  [30/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/common.cu.o
  [31/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_common_objects.dir/precision_independent.cu.o
  [32/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/1d/interp1d_wrapper.cu.o
  [33/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/3d/interp3d_wrapper.cu.o
  [34/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/2d/interp2d_wrapper.cu.o
  [35/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/1d/spread1d_wrapper.cu.o
  [36/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/2d/spread2d_wrapper.cu.o
  [37/42] Building CUDA object vendor/finufft/src/cuda/CMakeFiles/cufinufft_objects.dir/3d/spread3d_wrapper.cu.o
  [38/42] Linking CXX static library libcufinufft_static.a
  [39/42] Linking CXX shared library CMakeFiles/cufinufft.dir/cmake_device_link.o
  [40/42] Linking CXX shared library libcufinufft.so
  [41/42] Linking CXX shared module CMakeFiles/jax_finufft_gpu.dir/cmake_device_link.o
  [42/42] Linking CXX shared module jax_finufft_gpu.cpython-310-x86_64-linux-gnu.so
  *** Installing project into wheel...
  -- Install configuration: "Release"
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/lib/libfinufft.so
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/include/finufft.h
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/include/finufft_eitherprec.h
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/include/finufft_errors.h
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/include/finufft_opts.h
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/include/finufft_spread_opts.h
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/lib/libfinufft_static.a
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/lib/libcufinufft.so
  -- Set runtime path of "/tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/lib/libcufinufft.so" to ""
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/lib/libcufinufft_static.a
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/licenses/finufft/LICENSE
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/guru1d1c.c
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/simple1d1.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/threadsafe1d1.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/guru1d1f.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/simple1d1c.c
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/simple1d1cf.c
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/gurumany1d1.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/many1d1.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/simulplans1d1.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/guru2d1.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/guru1d1.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/simple2d1.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/threadsafe2d2f.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/simple1d1f.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/cuda
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/cuda/example2d1many.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/cuda/example2d2many.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/share/finufft/examples/cuda/getting_started.cpp
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/./jax_finufft_cpu.cpython-310-x86_64-linux-gnu.so
  -- Installing: /tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/./jax_finufft_gpu.cpython-310-x86_64-linux-gnu.so
  -- Set runtime path of "/tmp/tmp4qbez2ec/wheel/platlib/jax_finufft/./jax_finufft_gpu.cpython-310-x86_64-linux-gnu.so" to ""
  *** Making wheel...
  *** Created jax_finufft-0.0.4.dev53+g049f0b9-cp310-cp310-manylinux_2_35_x86_64.whl...
  Building wheel for jax-finufft (pyproject.toml) ... done
  Created wheel for jax-finufft: filename=jax_finufft-0.0.4.dev53+g049f0b9-cp310-cp310-manylinux_2_35_x86_64.whl size=104259500 sha256=05e3045b3c2c13f07f4fffcb9bd03c2907284eae3ba6783a98cf38e55675b813
  Stored in directory: /home/aparsons/.cache/pip/wheels/b3/09/ff/f4516de6307b667b70e4cc44ff10332554e970ad70dbe14a36
Successfully built jax-finufft
Installing collected packages: jax-finufft
  Attempting uninstall: jax-finufft
    Found existing installation: jax-finufft 0.0.4.dev53+g049f0b9
    Uninstalling jax-finufft-0.0.4.dev53+g049f0b9:
      Removing file or directory /home/aparsons/.local/lib/python3.10/site-packages/jax_finufft-0.0.4.dev53+g049f0b9.dist-info/
      Removing file or directory /home/aparsons/.local/lib/python3.10/site-packages/jax_finufft/
      Successfully uninstalled jax-finufft-0.0.4.dev53+g049f0b9
Successfully installed jax-finufft-0.0.4.dev53+g049f0b9




$ pytest ops_test.py::test_nufft1_forward[2-False-50-75--1]
============================= test session starts ==============================
platform linux -- Python 3.10.12, pytest-7.4.2, pluggy-1.3.0
rootdir: /home/aparsons/projects/hera/jax-finufft/tests
plugins: anyio-3.6.2
collected 1 item                                                               

ops_test.py Fatal Python error: Segmentation fault

Current thread 0x00007fa2bd4b9740 (most recent call first):
  File "/home/aparsons/.local/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py", line 1916 in __call__
  File "/home/aparsons/.local/lib/python3.10/site-packages/jax/_src/profiler.py", line 314 in wrapper
  File "/home/aparsons/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 1252 in _pjit_call_impl
  File "/home/aparsons/.local/lib/python3.10/site-packages/jax/_src/core.py", line 817 in process_primitive
  File "/home/aparsons/.local/lib/python3.10/site-packages/jax/_src/core.py", line 363 in bind_with_trace
  File "/home/aparsons/.local/lib/python3.10/site-packages/jax/_src/core.py", line 2592 in bind
  File "/home/aparsons/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 185 in _python_pjit_helper
  File "/home/aparsons/.local/lib/python3.10/site-packages/jax/_src/pjit.py", line 238 in cache_miss
  File "/home/aparsons/.local/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 166 in reraise_with_filtered_traceback
  File "/home/aparsons/projects/hera/jax-finufft/tests/ops_test.py", line 45 in test_nufft1_forward
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/python.py", line 194 in pytest_pyfunc_call
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_callers.py", line 77 in _multicall
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_manager.py", line 115 in _hookexec
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_hooks.py", line 493 in __call__
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/python.py", line 1792 in runtest
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/runner.py", line 169 in pytest_runtest_call
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_callers.py", line 77 in _multicall
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_manager.py", line 115 in _hookexec
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_hooks.py", line 493 in __call__
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/runner.py", line 262 in <lambda>
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/runner.py", line 341 in from_call
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/runner.py", line 261 in call_runtest_hook
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/runner.py", line 222 in call_and_report
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/runner.py", line 133 in runtestprotocol
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/runner.py", line 114 in pytest_runtest_protocol
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_callers.py", line 77 in _multicall
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_manager.py", line 115 in _hookexec
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_hooks.py", line 493 in __call__
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/main.py", line 350 in pytest_runtestloop
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_callers.py", line 77 in _multicall
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_manager.py", line 115 in _hookexec
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_hooks.py", line 493 in __call__
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/main.py", line 325 in _main
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/main.py", line 271 in wrap_session
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/main.py", line 318 in pytest_cmdline_main
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_callers.py", line 77 in _multicall
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_manager.py", line 115 in _hookexec
  File "/home/aparsons/.local/lib/python3.10/site-packages/pluggy/_hooks.py", line 493 in __call__
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/config/__init__.py", line 169 in main
  File "/home/aparsons/.local/lib/python3.10/site-packages/_pytest/config/__init__.py", line 192 in console_main
  File "/home/aparsons/.local/bin/pytest", line 8 in <module>

Extension modules: jaxlib.cpu_feature_guard, numpy.core._multiarray_umath, numpy.core._multiarray_tests, numpy.linalg._umath_linalg, numpy.fft._pocketfft_internal, numpy.random._common, numpy.random.bit_generator, numpy.random._bounded_integers, numpy.random._mt19937, numpy.random.mtrand, numpy.random._philox, numpy.random._pcg64, numpy.random._sfc64, numpy.random._generator (total: 14)
Segmentation fault (core dumped)

Actually, that segfault may have been from a stale GPU interface held by a notebook while running the tests in the command line. They have cleared, so I think we can call this issue resolved.

Here is the final list of instructions:
$ cd jax-finufft
$ nvidia-smi --query-gpu=compute_cap --format=csv,noheader
7.5
$ export CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=75 -DJAX_FINUFFT_USE_CUDA=ON"
$ echo $LD_LIBRARY_PATH
/usr/local/cuda-12.3/lib64:/usr/local/lib:
$ rm -rf build
$ python -m pip install "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
$ pip install pycuda
$ pip install pybind11[global]
$ export CMAKE_PREFIX_PATH=/home/aparsons/.local/lib/python3.10/site-packages/pybind11/share/cmake/pybind11
$ pip install .

A major thing I missed is that cmake on all submodules is run within the pip process, and I think a source of my original problems was related to the "cmake" instructions that were linked, which implied that I needed to run some of the cmake build of submodules manually.
Thanks!
Thanks for reporting back, glad it's working!

> I think a source of my original problems was related to the "cmake" instructions that were linked, which implied that I needed to run some of the cmake build of submodules manually.

This makes sense, I can see how that would be confusing. I'll update the docs; really, there's not much reason for us to link to that page, we can just echo the instructions about querying `nvidia-smi`.

These three lines should definitely not be necessary, though:
>  $ pip install pycuda
> $ pip install pybind11[global]
> $ export CMAKE_PREFIX_PATH=/home/aparsons/.local/lib/python3.10/site-packages/pybind11/share/cmake/pybind11

If it wasn't working without these, it might be worth trying a Python venv next time to see if that helps. Probably not worth messing with your working installation right now, but just wanted to mention it in case it comes up again.

Fix patch:
diff --git a/.github/workflows/gpu-build.yml b/.github/workflows/gpu-build.yml
index 5739bcc..00fb6d8 100644
--- a/.github/workflows/gpu-build.yml
+++ b/.github/workflows/gpu-build.yml
@@ -30,7 +30,9 @@ jobs:
           python-version: 3.11
 
       - name: Install Python dependencies
-        run: python -m pip install -U pip
+        run: |
+          python -m pip install -U pip
+          python -m pip install -U jax[cpu]
 
       - name: Compile the extension
         run: python -m pip install -v .
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 099b558..2760de0 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -10,7 +10,7 @@ message(STATUS "Using CMake version: " ${CMAKE_VERSION})
 set(CMAKE_INTERPROCEDURAL_OPTIMIZATION OFF)
 
 # Enable CUDA if requested and available
-option(JAX_FINUFFT_USE_CUDA "Enable CUDA build if available" OFF)
+option(JAX_FINUFFT_USE_CUDA "Enable CUDA build" OFF)
 if(JAX_FINUFFT_USE_CUDA)
     include(CheckLanguage)
     check_language(CUDA)
@@ -19,7 +19,8 @@ if(JAX_FINUFFT_USE_CUDA)
         enable_language(CUDA)
         set(FINUFFT_USE_CUDA ON)
     else()
-        message(STATUS "No CUDA compiler found; compiling without GPU support")
+        message(FATAL_ERROR "No CUDA compiler found! Please ensure the CUDA Toolkit "
+            "is installed, or set JAX_FINUFFT_USE_CUDA=OFF to disable GPU support.")
         set(FINUFFT_USE_CUDA OFF)
     endif()
 else()
diff --git a/README.md b/README.md
index 032e27d..e7d288e 100644
--- a/README.md
+++ b/README.md
@@ -30,7 +30,7 @@ _For now, only a source build is supported._
 For building, you should only need a recent version of Python (>3.6) and
 [FFTW](https://www.fftw.org/). GPU-enabled builds also require a working CUDA
 compiler (i.e. the CUDA Toolkit), CUDA >= 11.8, and a compatible cuDNN (older versions of CUDA may work but
-are untested). At runtime, you'll need `numpy`, `scipy`, and `jax`.
+are untested). At runtime, you'll need `numpy` and `jax`.
 
 First, clone the repo and `cd` into the repo root (don't forget the `--recursive` flag because FINUFFT is included as a submodule):
 ```bash
@@ -52,7 +52,7 @@ python -m pip install .
 The `CPATH` export is needed so that the build can find the headers for libraries like FFTW installed through conda.
 
 For a GPU build, while the CUDA libraries and compiler are nominally available through conda,
-our experience trying to install them through conda suggests that the "traditional"
+our experience trying to install them this way suggests that the "traditional"
 way of obtaining the [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads) directly
 from NVIDIA may work best (see [related advice for Horovod](https://horovod.readthedocs.io/en/stable/conda_include.html)). After installing the CUDA Toolkit, one can set up the rest of the dependencies with:
 
@@ -65,7 +65,22 @@ python -m pip install "jax[cuda11_local]" -f https://storage.googleapis.com/jax-
 python -m pip install .
 ```
 
-In the `CMAKE_ARGS` line, you'll need to select the CUDA architecture(s) you wish to compile for. See the [FINUFFT docs](https://finufft.readthedocs.io/en/latest/install_gpu.html#cmake-installation).
+Other ways of installing JAX are given on the JAX website; the ["local CUDA" install methods](https://jax.readthedocs.io/en/latest/installation.html#pip-installation-gpu-cuda-installed-locally-harder) are preferred for jax-finufft as this ensures the CUDA extensions are compiled with the same Toolkit version as the CUDA runtime.
+
+In the above `CMAKE_ARGS` line, you'll need to select the CUDA architecture(s) you wish to compile for. To query your GPU's CUDA architecture (compute capability), you can run:
+
+```bash
+$ nvidia-smi --query-gpu=compute_cap --format=csv,noheader
+7.0
+```
+
+This corresponds to `CMAKE_CUDA_ARCHITECTURES=70`, i.e.:
+
+```bash
+export CMAKE_ARGS="-DCMAKE_CUDA_ARCHITECTURES=70 -DJAX_FINUFFT_USE_CUDA=ON"
+```
+
+Note that the pip installation is running CMake, so `CMAKE_ARGS` has to be set before then, but is not needed at runtime.
 
 At runtime, you may also need:
 ```bash
diff --git a/pyproject.toml b/pyproject.toml
index 5ca9f74..d467e44 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ authors = [{ name = "Dan Foreman-Mackey", email = "foreman.mackey@gmail.com" }]
 requires-python = ">=3.8"
 license = { file = "LICENSE" }
 urls = { Homepage = "https://github.com/dfm/jax-finufft" }
-dependencies = ["jax", "jaxlib"]
+dependencies = ["jax", "numpy"]
 dynamic = ["version"]
 
 [project.optional-dependencies]


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 0492834..47a858b 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -37,6 +37,7 @@ jobs:
       - name: Build
         run: |
           python -m pip install -U pip
+          python -m pip install -U jax[cpu]
           python -m pip install -v .[test]
 
       - name: Run tests