Address code reviews

Co-authored-by: Diego Ferigo <diego.ferigo@iit.it>

src/jaxsim/__init__.py

@@ -14,13 +14,14 @@ def _jnp_options() -> None:
     # Determine if 64-bit precision is requested
     use_x64 = os.environ.get("JAX_ENABLE_X64", "1") != "0"
 
-    # Raise an error if 64-bit precision is not allowed on TPU
+    # Notify the user if unsupported 64-bit precision was enforced on TPU.
     if is_tpu and use_x64:
-        msg = "64-bit precision is not allowed on TPU. Enforcing 32bit precision."
-        logging.error(msg)
+        msg = "64bit precision is not allowed on TPU. Enforcing 32bit precision."
+        logging.warning(msg)
+        use_x64 = False
 
     # Enable 64-bit precision in JAX
-    elif not is_tpu and use_x64:
+    if use_x64:
         logging.info("Enabling JAX to use 64bit precision")
         jax.config.update("jax_enable_x64", True)
 
@@ -31,10 +32,6 @@ def _jnp_options() -> None:
         if jnp.empty(0, dtype=float).dtype != jnp.empty(0, dtype=np.float64).dtype:
             logging.warning("Failed to enable 64bit precision in JAX")
 
-    # Enforce 32-bit precision on TPU
-    elif is_tpu:
-        logging.warning("JAX is running on TPU; 32bit precision is enforced.")
-
     # Warn about experimental use of 32-bit precision
     else:
         logging.warning(