don't include ffi headers if cuda is not available

Author: ASKabalan

lib/src/extensions.cc

@@ -1,21 +1,22 @@
 #include <nanobind/nanobind.h>
-#include "xla/ffi/api/api.h"
-#include "xla/ffi/api/c_api.h"
-#include "xla/ffi/api/ffi.h"
 #include <cstddef>
 #include <complex>
 #include <type_traits>
 
-namespace ffi = xla::ffi;
 namespace nb = nanobind;
 
 #ifndef NO_CUDA_COMPILER
+#include "xla/ffi/api/api.h"
+#include "xla/ffi/api/c_api.h"
+#include "xla/ffi/api/ffi.h"
 #include "cuda_runtime.h"
 #include "plan_cache.h"
 #include "s2fft_kernels.h"
 #include "s2fft.h"
 #include "cudastreamhandler.hpp"  // For forking and joining CUDA streams
 
+namespace ffi = xla::ffi;
+
 namespace s2fft {
 
 /**

s2fft/transforms/spherical.py

@@ -261,7 +261,7 @@ def inverse_jax(
             "healpix"` and running on a CUDA compatible GPU device. Using a custom
             primitive reduces long compilation times when just-in-time compiling.
             Defaults to `False`.
-Z
+
     Returns:
         jnp.ndarray: Signal on the sphere.