import jaxns only if double precision is enabled

pyro-ppl · Oct 16, 2024 · 19ae8f0 · 19ae8f0
1 parent 95937ff
commit 19ae8f0
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/test/contrib/test_nested_sampling.py b/test/contrib/test_nested_sampling.py
@@ -14,12 +14,14 @@
 import numpyro
 
 try:
-    with warnings.catch_warnings():
-        warnings.filterwarnings("ignore", category=UserWarning)
-        from numpyro.contrib.nested_sampling import NestedSampler, UniformReparam
+    if os.environ.get("JAX_ENABLE_X64"):
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=UserWarning)
+            from numpyro.contrib.nested_sampling import NestedSampler, UniformReparam
 
 except ImportError:
     pytestmark = pytest.mark.skip(reason="jaxns is not installed")
+
 import numpyro.distributions as dist
 from numpyro.distributions.transforms import AffineTransform, ExpTransform