changed default for float32 to float64 with jax

Jashcraf · Oct 11, 2024 · 165bfb0 · 165bfb0
1 parent f1e36cc
commit 165bfb0
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/katsu/katsu_math.py b/katsu/katsu_math.py
@@ -38,9 +38,10 @@ def set_backend_to_cupy():
 
 def set_backend_to_jax(): 
     """Convenience method to automatically configure katsu's backend to jax."""
-    import jax.numpy as jnp
+    import jax as jax
 
-    np._srcmodule = jnp
+    jax.config.update("jax_enable_x64", True)
+    np._srcmodule = jax.numpy
 
     return