jaxsim 0.4.3.dev177__py3-none-any.whl → 0.4.3.dev181__py3-none-any.whl

jaxsim/__init__.py

@@ -8,21 +8,35 @@ def _jnp_options() -> None:
 
     import jax
 
-    # Enable by default 64bit precision in JAX.
-    if os.environ.get("JAX_ENABLE_X64", "1") != "0":
-
-        logging.info("Enabling JAX to use 64bit precision")
+    # Check if running on TPU
+    is_tpu = jax.devices()[0].platform == "tpu"
+
+    # Enable by default 64-bit precision to get accurate physics.
+    # Users can enforce 32-bit precision by setting the following variable to 0.
+    use_x64 = os.environ.get("JAX_ENABLE_X64", "1") != "0"
+
+    # 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.warning(msg)
+        use_x64 = False
+
+    # Enable 64-bit precision in JAX.
+    if use_x64:
+        logging.info("Enabling JAX to use 64-bit precision")
         jax.config.update("jax_enable_x64", True)
 
         import jax.numpy as jnp
         import numpy as np
 
+        # Verify that 64-bit precision is correctly set.
         if jnp.empty(0, dtype=float).dtype != jnp.empty(0, dtype=np.float64).dtype:
-            logging.warning("Failed to enable 64bit precision in JAX")
+            logging.warning("Failed to enable 64-bit precision in JAX")
 
+    # Warn about experimental usage of 32-bit precision.
     else:
         logging.warning(
-            "Using 32bit precision in JaxSim is still experimental, please avoid to use variable step integrators."
+            "Using 32-bit precision in JaxSim is still experimental, please avoid to use variable step integrators."
         )
 
 

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev177'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev177')
+__version__ = version = '0.4.3.dev181'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev181')

jaxsim/exceptions.py

@@ -17,6 +17,10 @@ def raise_if(
             format string (fmt), whose fields are filled with the args and kwargs.
     """
 
+    # Disable host callback if running on TPU.
+    if jax.devices()[0].platform == "tpu":
+        return
+
     # Check early that the format string is well-formed.
     try:
         _ = msg.format(*args, **kwargs)

{jaxsim-0.4.3.dev177.dist-info → jaxsim-0.4.3.dev181.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev177
+Version: 0.4.3.dev181
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.3.dev177.dist-info → jaxsim-0.4.3.dev181.dist-info}/RECORD

@@ -1,6 +1,6 @@
-jaxsim/__init__.py,sha256=bSbpggIz5aG6QuGZLa0V2EfHjAOeucMxi-vIYxzLmN8,2788
-jaxsim/_version.py,sha256=SzGoIDpeznpZHWyfdxXEtnO3y8zLaXJORHRJRSUcxsU,428
-jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
+jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
+jaxsim/_version.py,sha256=2YMpT461ObsI3rceAHrUVR-OPYPsQ43SHkzRH20mlDY,428
+jaxsim/exceptions.py,sha256=vSoScaRD4nvh6jltgK9Ry5pKnE0O5hb4_yI_pk_fvR8,2175
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
@@ -64,8 +64,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev177.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev177.dist-info/METADATA,sha256=BMlT_szB4WbLIZmCucz750x_aRXoe2n7ycwlzI3l-sk,17276
-jaxsim-0.4.3.dev177.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-jaxsim-0.4.3.dev177.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev177.dist-info/RECORD,,
+jaxsim-0.4.3.dev181.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev181.dist-info/METADATA,sha256=w-b474j6ugFST4V5QJ2h9DTtrb2xRH0cdCRAMyBG8wg,17276
+jaxsim-0.4.3.dev181.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+jaxsim-0.4.3.dev181.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev181.dist-info/RECORD,,