missing conditions in setup, segfault test

setup.py

@@ -42,6 +42,10 @@ def run(self):
             f"-DSPHERICART_ARCH_NATIVE={SPHERICART_ARCH_NATIVE}",
         ]
 
+        CUDA_HOME = os.environ.get("CUDA_HOME")
+        if CUDA_HOME is not None:
+            cmake_options.append(f"-DCUDA_TOOLKIT_ROOT_DIR={CUDA_HOME}")
+
         if sys.platform.startswith("darwin"):
             cmake_options.append("-DCMAKE_OSX_DEPLOYMENT_TARGET:STRING=11.0")
 

sphericart-jax/python/tests/segfault_test.py

@@ -0,0 +1,22 @@
+import pytest
+import jax
+
+# jax.config.update("jax_platform_name", "cpu")
+import jax.numpy as jnp
+import jax._src.test_util as jtu
+import sphericart.jax
+
+def xyz():
+    key = jax.random.PRNGKey(0)
+    return 6 * jax.random.normal(key, (100, 3))
+
+
+def compute(xyz):
+    sph = sphericart.jax.spherical_harmonics(l_max=4, normalized=False, xyz=xyz)
+    assert jnp.linalg.norm(sph) != 0.0
+    return sph.sum()
+
+
+sph = compute(xyz())
+
+print (sph)

sphericart-jax/setup.py

@@ -30,9 +30,12 @@ def run(self):
             f"-DPYTHON_EXECUTABLE={sys.executable}",
             f"-DSPHERICART_ARCH_NATIVE={SPHERICART_ARCH_NATIVE}",
             f"-DCMAKE_PREFIX_PATH={';'.join(cmake_prefix_path)}",
-            "-DSPHERICART_ENABLE_CUDA=ON"
         ]
 
+        CUDA_HOME = os.environ.get("CUDA_HOME")
+        if CUDA_HOME is not None:
+            cmake_options.append(f"-DCUDA_TOOLKIT_ROOT_DIR={CUDA_HOME}")
+
         if sys.platform.startswith("darwin"):
             cmake_options.append("-DCMAKE_OSX_DEPLOYMENT_TARGET:STRING=11.0")