search

lab-cosmo / sphericart

Multi-language library for the calculation of spherical harmonics in Cartesian coordinates
https://sphericart.readthedocs.io/en/latest/
MIT License
73 stars 13 forks source link

Installing sphericart.jax with --no-build-isolation in an environment containing jax[cuda12] (0.4.34) yields XLA failures #151

Open nickjbrowning opened 1 month ago

nickjbrowning commented 1 month ago

Only seems to occur on hessians.

python -m pytest sphericart-jax/python/tests/test_transforms.py

FAILED sphericart-jax/python/tests/test_transforms.py::test_vmap_hessian[4-True] - jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the library was not initialized FAILED sphericart-jax/python/tests/test_transforms.py::test_vmap_hessian[4-False] - jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the library was not initialized FAILED sphericart-jax/python/tests/test_transforms.py::test_vmap_hessian[7-True] - jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the library was not initialized FAILED sphericart-jax/python/tests/test_transforms.py::test_vmap_hessian[7-False] - jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the library was not initialized FAILED sphericart-jax/python/tests/test_transforms.py::test_vmap_hessian[10-True] - jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the library was not initialized FAILED sphericart-jax/python/tests/test_transforms.py::test_vmap_hessian[10-False] - jaxlib.xla_extension.XlaRuntimeError: INTERNAL: the library was not initialized