Update Pyproject.toml and build to include FFI headers

Author: ASKabalan

CMakeLists.txt

@@ -32,8 +32,11 @@ if(CMAKE_CUDA_COMPILER)
     find_package(Python 3.8
     REQUIRED COMPONENTS Interpreter Development.Module
     OPTIONAL_COMPONENTS Development.SABIModule)
-    set(XLA_DIR ${Python_SITELIB}/jaxlib/include)
-    message(STATUS "XLA_DIR: ${XLA_DIR}")
+    execute_process(
+    COMMAND "${Python_EXECUTABLE}"
+            "-c" "from jax.extend import ffi; print(ffi.include_dir())"
+    OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE XLA_DIR)
+    message(STATUS "XLA include directory: ${XLA_DIR}")
 
     # Detect the installed nanobind package and import it into CMake
     find_package(nanobind CONFIG REQUIRED)

pyproject.toml

@@ -2,8 +2,9 @@
 requires = [
     "setuptools",
     "setuptools-scm",
-    "scikit-build-core >=0.4.3", 
-    "nanobind >=1.3.2"
+    "scikit-build-core >=0.11", 
+    "nanobind >=2.0,<2.6",
+    "jax >= 0.4.0"
 ]
 build-backend = "scikit_build_core.build"
 
@@ -78,11 +79,12 @@ tests = [
 
 [tool.scikit-build]
 # Protect the configuration against future changes in scikit-build-core
-minimum-version = "0.4"
+minimum-version = "0.8"
 # Setuptools-style build caching in a local directory
 build-dir = "build/{wheel_tag}"
 # Build stable ABI wheels for CPython 3.12+
 wheel.py-api = "cp312"
+cmake.build-type = "Release"
 metadata.version.provider = "scikit_build_core.metadata.setuptools_scm"
 sdist.include = ["s2fft/_version.py"]