Merge pull request #92 from ami-iit/flferretti-patch-1

ami-iit · Feb 28, 2024 · 6a5040b · 6a5040b
2 parents 389c556 + b5b4f31
commit 6a5040b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/jaxsim/__init__.py b/src/jaxsim/__init__.py
@@ -6,12 +6,12 @@
 def _jnp_options() -> None:
     import os
 
-    from jax.config import config
+    import jax
 
     # Enable by default
     if not ("JAX_ENABLE_X64" in os.environ and os.environ["JAX_ENABLE_X64"] == "0"):
         logging.info("Enabling JAX to use 64bit precision")
-        config.update("jax_enable_x64", True)
+        jax.config.update("jax_enable_x64", True)
 
         import jax.numpy as jnp
         import numpy as np