jaxsim 0.5.1.dev95__py3-none-any.whl → 0.5.1.dev103__py3-none-any.whl

jaxsim/__init__.py

@@ -8,19 +8,27 @@ def _jnp_options() -> None:
 
     import jax
 
-    # Check if running on TPU
+    # Check if running on TPU.
     is_tpu = jax.devices()[0].platform == "tpu"
 
+    # Check if running on Metal.
+    is_metal = jax.devices()[0].platform == "METAL"
+
     # 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."
+    if (is_tpu or is_metal) and use_x64:
+        msg = f"64-bit precision is not allowed on {jax.devices()[0].platform.upper}. Enforcing 32bit precision."
         logging.warning(msg)
         use_x64 = False
 
+        if is_metal:
+            logging.warning(
+                "JAX Metal backend is experimental. Some functionalities may not be available."
+            )
+
     # Enable 64-bit precision in JAX.
     if use_x64:
         logging.info("Enabling JAX to use 64-bit precision")

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.1.dev95'
-__version_tuple__ = version_tuple = (0, 5, 1, 'dev95')
+__version__ = version = '0.5.1.dev103'
+__version_tuple__ = version_tuple = (0, 5, 1, 'dev103')

jaxsim/exceptions.py

@@ -19,8 +19,9 @@ 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" or os.environ.get(
+    # Disable host callback if running on unsupported hardware or if the user
+    # explicitly disabled it.
+    if jax.devices()[0].platform in {"tpu", "METAL"} or os.environ.get(
         "JAXSIM_DISABLE_EXCEPTIONS", 0
     ):
         return

jaxsim/parsers/rod/utils.py

@@ -225,15 +225,15 @@ def create_mesh_collision(
 ) -> descriptions.MeshCollision:
 
     file = pathlib.Path(resolve_local_uri(uri=collision.geometry.mesh.uri))
-    _file_type = file.suffix.replace(".", "")
-    mesh = trimesh.load_mesh(file, file_type=_file_type)
+    file_type = file.suffix.replace(".", "")
+    mesh = trimesh.load_mesh(file, file_type=file_type)
 
     if mesh.is_empty:
         raise RuntimeError(f"Failed to process '{file}' with trimesh")
 
     mesh.apply_scale(collision.geometry.mesh.scale)
     logging.info(
-        msg=f"Loading mesh {collision.geometry.mesh.uri} with scale {collision.geometry.mesh.scale}, file type '{_file_type}'"
+        msg=f"Loading mesh {collision.geometry.mesh.uri} with scale {collision.geometry.mesh.scale}, file type '{file_type}'"
     )
 
     if method is None:

jaxsim/typing.py

@@ -20,9 +20,9 @@ PyTree: object = (
     dict[Hashable, TypeVar("PyTree")]
     | list[TypeVar("PyTree")]
     | tuple[TypeVar("PyTree")]
-    | None
     | jax.Array
     | Any
+    | None
 )
 
 # =======================

{jaxsim-0.5.1.dev95.dist-info → jaxsim-0.5.1.dev103.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev95
+Version: 0.5.1.dev103
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.5.1.dev95.dist-info → jaxsim-0.5.1.dev103.dist-info}/RECORD

@@ -1,8 +1,8 @@
-jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
-jaxsim/_version.py,sha256=uLmioR8f_16hLO76GPSpzQyclhE2S-4nH8pntMObYGY,426
-jaxsim/exceptions.py,sha256=Sq3qtqeiy-CK76new_W2KKQ-4MAzyOUK5j5pBLr4RPQ,2250
+jaxsim/__init__.py,sha256=OQcCxXn4BXiSvGjAatqvw5fAEVfXbxyavQZN25NEyBo,3675
+jaxsim/_version.py,sha256=zdHXFLMTP9gf9Mb93H_-kwqd5cnJDwjqSFX1iUItBJ4,428
+jaxsim/exceptions.py,sha256=5rYd0kGrIaq1e0WK0qMVTCkoGJtMWz_uWlVtxHbjH_A,2322
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
-jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
+jaxsim/typing.py,sha256=7msl8t5Jt09RNYfKdPJtpjLfWurldcycDappb045Eso,761
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=5fYNRUhKE5VGGdW88zY8mqqEy5VTWyaHu5k6MgW4Jt4,13826
 jaxsim/api/common.py,sha256=SvEOGxCKOxKLLVHaNp1sFkBX0sku3-wH0-HUlYVWCDk,7090
@@ -46,7 +46,7 @@ jaxsim/parsers/descriptions/model.py,sha256=I2Vsbv8Josl4Le7b5rIvhqA2k9Bbv5JxMqwy
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
 jaxsim/parsers/rod/meshes.py,sha256=a5qUFkHe3gPXMKRnvPx3BDvlEHT2vEJqcLYLXhJFCdA,2847
 jaxsim/parsers/rod/parser.py,sha256=OJdKM2hVyED9nw7iRpxCRFZx9Z6rKmVflhkItGTzRns,14499
-jaxsim/parsers/rod/utils.py,sha256=7qKiTQk897f_awKt1Dd7OGTx6zUPvv7r5yYODMzXEaY,7833
+jaxsim/parsers/rod/utils.py,sha256=qvaWQKbmUWu1t-IOT2eseBHCMm3tdBQ1XuYSBO6IDXY,7830
 jaxsim/rbda/__init__.py,sha256=kmy4G9aMkrqPNGdLSaSV3k15dpF52vBEUQXDFDuKIxU,337
 jaxsim/rbda/aba.py,sha256=w7ciyxB0IsmueatT0C7PcBQEl9dyiH9oqJgIi3xeTUE,8983
 jaxsim/rbda/collidable_points.py,sha256=0PFLzxWKtRg8-JtfNhGlSjBMv1J98tiLymOdvlvAak4,5325
@@ -67,8 +67,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=eEY28MZW0Lm_jJNt1NkFqZz0ek01tvhR46OXZYCo7tc,532
 jaxsim/utils/wrappers.py,sha256=ZY7olSORzZRvSzkdeNLj8yjwUIAt9L0Douwl7wItjpk,4008
-jaxsim-0.5.1.dev95.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.5.1.dev95.dist-info/METADATA,sha256=xOMthPKBfJqlJrzyVADzPjB5b-h8MKI3K8CGRdcKZy4,17937
-jaxsim-0.5.1.dev95.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
-jaxsim-0.5.1.dev95.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.5.1.dev95.dist-info/RECORD,,
+jaxsim-0.5.1.dev103.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.5.1.dev103.dist-info/METADATA,sha256=QWMkkVWOPO0aedePnBihvUEV395YxLafSFGzxnYVmVE,17938
+jaxsim-0.5.1.dev103.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
+jaxsim-0.5.1.dev103.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.5.1.dev103.dist-info/RECORD,,