jaxsim 0.2.dev2__tar.gz → 0.2.dev8__tar.gz

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.dev2
+Version: 0.2.dev8
 Summary: A physics engine in reduced coordinates implemented with JAX.
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
@@ -38,7 +38,7 @@ Requires-Dist: jax_dataclasses>=1.4.0
 Requires-Dist: pptree
 Requires-Dist: rod
 Provides-Extra: style
-Requires-Dist: black; extra == "style"
+Requires-Dist: black[jupyter]; extra == "style"
 Requires-Dist: isort; extra == "style"
 Provides-Extra: testing
 Requires-Dist: idyntree; extra == "testing"
@@ -47,7 +47,7 @@ Requires-Dist: pytest-forked; extra == "testing"
 Requires-Dist: pytest-icdiff; extra == "testing"
 Requires-Dist: robot-descriptions; extra == "testing"
 Provides-Extra: all
-Requires-Dist: black; extra == "all"
+Requires-Dist: black[jupyter]; extra == "all"
 Requires-Dist: isort; extra == "all"
 Requires-Dist: idyntree; extra == "all"
 Requires-Dist: pytest>=6.0; extra == "all"

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/examples/PD_controller.ipynb

@@ -19,7 +19,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#@title Imports and setup\n",
+    "# @title Imports and setup\n",
     "from IPython.display import clear_output, HTML, display\n",
     "import sys\n",
     "\n",
@@ -52,7 +52,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#@title Fetch the URDF file\n",
+    "# @title Fetch the URDF file\n",
     "import requests\n",
     "\n",
     "url = \"https://raw.githubusercontent.com/ami-iit/jaxsim/main/examples/assets/cartpole.urdf\"\n",
@@ -117,7 +117,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#@title Set up MuJoCo renderer\n",
+    "# @title Set up MuJoCo renderer\n",
     "!{sys.executable} -m pip install -U -q mujoco\n",
     "!{sys.executable} -m pip install -q mediapy\n",
     "\n",
@@ -131,96 +131,106 @@
     "import subprocess\n",
     "\n",
     "if IS_COLAB:\n",
-    "  if subprocess.run('ffmpeg -version', shell=True).returncode:\n",
-    "    !command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
-    "    clear_output()\n",
+    "    if subprocess.run(\"ffmpeg -version\", shell=True).returncode:\n",
+    "        !command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
+    "        clear_output()\n",
+    "\n",
+    "    if subprocess.run(\"nvidia-smi\").returncode:\n",
+    "        raise RuntimeError(\n",
+    "            \"Cannot communicate with GPU. \"\n",
+    "            \"Make sure you are using a GPU Colab runtime. \"\n",
+    "            \"Go to the Runtime menu and select Choose runtime type.\"\n",
+    "        )\n",
     "\n",
-    "  if subprocess.run('nvidia-smi').returncode:\n",
-    "    raise RuntimeError(\n",
-    "        'Cannot communicate with GPU. '\n",
-    "        'Make sure you are using a GPU Colab runtime. '\n",
-    "        'Go to the Runtime menu and select Choose runtime type.')\n",
-    "\n",
-    "# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n",
-    "# This is usually installed as part of an Nvidia driver package, but the Colab\n",
-    "# kernel doesn't install its driver via APT, and as a result the ICD is missing.\n",
-    "# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n",
-    "  NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'\n",
-    "  if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n",
-    "    with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:\n",
-    "      f.write(\"\"\"{\n",
+    "    # Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n",
+    "    # This is usually installed as part of an Nvidia driver package, but the Colab\n",
+    "    # kernel doesn't install its driver via APT, and as a result the ICD is missing.\n",
+    "    # (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n",
+    "    NVIDIA_ICD_CONFIG_PATH = \"/usr/share/glvnd/egl_vendor.d/10_nvidia.json\"\n",
+    "    if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n",
+    "        with open(NVIDIA_ICD_CONFIG_PATH, \"w\") as f:\n",
+    "            f.write(\n",
+    "                \"\"\"{\n",
     "      \"file_format_version\" : \"1.0.0\",\n",
     "      \"ICD\" : {\n",
     "         \"library_path\" : \"libEGL_nvidia.so.0\"\n",
     "     }\n",
     "  }\n",
-    "  \"\"\")\n",
+    "  \"\"\"\n",
+    "            )\n",
     "\n",
     "%env MUJOCO_GL=egl\n",
     "\n",
     "try:\n",
-    "  import mujoco\n",
+    "    import mujoco\n",
     "except Exception as e:\n",
-    "  raise e from RuntimeError(\n",
-    "      'Something went wrong during installation. Check the shell output above '\n",
-    "      'for more information.\\n'\n",
-    "      'If using a hosted Colab runtime, make sure you enable GPU acceleration '\n",
-    "      'by going to the Runtime menu and selecting \"Choose runtime type\".')\n",
+    "    raise e from RuntimeError(\n",
+    "        \"Something went wrong during installation. Check the shell output above \"\n",
+    "        \"for more information.\\n\"\n",
+    "        \"If using a hosted Colab runtime, make sure you enable GPU acceleration \"\n",
+    "        'by going to the Runtime menu and selecting \"Choose runtime type\".'\n",
+    "    )\n",
+    "\n",
     "\n",
     "def load_mujoco_model_with_camera(xml_string, camera_pos, camera_xyaxes):\n",
-    "  def to_mjcf_string(list_to_str):\n",
-    "    return ' '.join(map(str, list_to_str))\n",
-    "\n",
-    "  mj_model_raw = mujoco.MjModel.from_xml_string(model_urdf_string)\n",
-    "  path_temp_xml = tempfile.NamedTemporaryFile(mode=\"w+\")\n",
-    "  mujoco.mj_saveLastXML(path_temp_xml.name, mj_model_raw)\n",
-    "  # Add camera in mujoco model\n",
-    "  tree = ET.parse(path_temp_xml)\n",
-    "  for elem in tree.getroot().iter(\"worldbody\"):\n",
-    "    worldbody_elem = elem\n",
-    "  camera_elem = ET.Element(\"camera\")\n",
-    "  # Set attributes  \n",
-    "  camera_elem.set(\"name\", \"side\")\n",
-    "  camera_elem.set(\"pos\", to_mjcf_string(camera_pos))\n",
-    "  camera_elem.set(\"xyaxes\", to_mjcf_string(camera_xyaxes))\n",
-    "  camera_elem.set(\"mode\", \"fixed\")\n",
-    "  worldbody_elem.append(camera_elem)\n",
-    "\n",
-    "  # Save new model\n",
-    "  mujoco_xml_with_camera = ET.tostring(tree.getroot(), encoding=\"unicode\")\n",
-    "  mj_model = mujoco.MjModel.from_xml_string(mujoco_xml_with_camera)\n",
-    "  return mj_model\n",
+    "    def to_mjcf_string(list_to_str):\n",
+    "        return \" \".join(map(str, list_to_str))\n",
+    "\n",
+    "    mj_model_raw = mujoco.MjModel.from_xml_string(model_urdf_string)\n",
+    "    path_temp_xml = tempfile.NamedTemporaryFile(mode=\"w+\")\n",
+    "    mujoco.mj_saveLastXML(path_temp_xml.name, mj_model_raw)\n",
+    "    # Add camera in mujoco model\n",
+    "    tree = ET.parse(path_temp_xml)\n",
+    "    for elem in tree.getroot().iter(\"worldbody\"):\n",
+    "        worldbody_elem = elem\n",
+    "    camera_elem = ET.Element(\"camera\")\n",
+    "    # Set attributes\n",
+    "    camera_elem.set(\"name\", \"side\")\n",
+    "    camera_elem.set(\"pos\", to_mjcf_string(camera_pos))\n",
+    "    camera_elem.set(\"xyaxes\", to_mjcf_string(camera_xyaxes))\n",
+    "    camera_elem.set(\"mode\", \"fixed\")\n",
+    "    worldbody_elem.append(camera_elem)\n",
+    "\n",
+    "    # Save new model\n",
+    "    mujoco_xml_with_camera = ET.tostring(tree.getroot(), encoding=\"unicode\")\n",
+    "    mj_model = mujoco.MjModel.from_xml_string(mujoco_xml_with_camera)\n",
+    "    return mj_model\n",
     "\n",
     "\n",
     "def from_jaxsim_to_mujoco_pos(jaxsim_jointpos, mjmodel, jaxsimmodel):\n",
-    "  mujocoqposaddr2jaxindex = {}\n",
-    "  for jaxjnt in jaxsimmodel.joints():\n",
-    "    jntname = jaxjnt.name()\n",
-    "    mujocoqposaddr2jaxindex[mjmodel.joint(jntname).qposadr[0]] = jaxjnt.index()-1\n",
+    "    mujocoqposaddr2jaxindex = {}\n",
+    "    for jaxjnt in jaxsimmodel.joints():\n",
+    "        jntname = jaxjnt.name()\n",
+    "        mujocoqposaddr2jaxindex[mjmodel.joint(jntname).qposadr[0]] = jaxjnt.index() - 1\n",
     "\n",
-    "  mujoco_jointpos = jaxsim_jointpos\n",
-    "  for i in range(0, len(mujoco_jointpos)):\n",
-    "    mujoco_jointpos[i] = jaxsim_jointpos[mujocoqposaddr2jaxindex[i]]\n",
+    "    mujoco_jointpos = jaxsim_jointpos\n",
+    "    for i in range(0, len(mujoco_jointpos)):\n",
+    "        mujoco_jointpos[i] = jaxsim_jointpos[mujocoqposaddr2jaxindex[i]]\n",
+    "\n",
+    "    return mujoco_jointpos\n",
     "\n",
-    "  return mujoco_jointpos\n",
-    " \n",
     "\n",
     "# To get a good camera location, you can use \"Copy camera\" functionality in MuJoCo GUI\n",
-    "mj_model = load_mujoco_model_with_camera(model_urdf_string, [3.954, 3.533, 2.343], [-0.594, 0.804, -0.000, -0.163, -0.120, 0.979])\n",
+    "mj_model = load_mujoco_model_with_camera(\n",
+    "    model_urdf_string,\n",
+    "    [3.954, 3.533, 2.343],\n",
+    "    [-0.594, 0.804, -0.000, -0.163, -0.120, 0.979],\n",
+    ")\n",
     "renderer = mujoco.Renderer(mj_model, height=480, width=640)\n",
     "\n",
+    "\n",
     "def get_image(camera, mujocojointpos) -> np.ndarray:\n",
-    "  \"\"\"Renders the environment state.\"\"\"\n",
-    "  # Copy joint data in mjdata state\n",
-    "  d = mujoco.MjData(mj_model)\n",
-    "  d.qpos = mujocojointpos\n",
-    "    \n",
-    "  # Forward kinematics\n",
-    "  mujoco.mj_forward(mj_model, d)\n",
-    "    \n",
-    "  # use the mjData object to update the renderer\n",
-    "  renderer.update_scene(d, camera=camera)\n",
-    "  return renderer.render()"
+    "    \"\"\"Renders the environment state.\"\"\"\n",
+    "    # Copy joint data in mjdata state\n",
+    "    d = mujoco.MjData(mj_model)\n",
+    "    d.qpos = mujocojointpos\n",
+    "\n",
+    "    # Forward kinematics\n",
+    "    mujoco.mj_forward(mj_model, d)\n",
+    "\n",
+    "    # use the mjData object to update the renderer\n",
+    "    renderer.update_scene(d, camera=camera)\n",
+    "    return renderer.render()"
    ]
   },
   {
@@ -241,10 +251,19 @@
     "sim_images = []\n",
     "timestep = 0.01\n",
     "for _ in range(300):\n",
-    "    sim_images.append(get_image(\"side\", from_jaxsim_to_mujoco_pos(np.array(model.joint_positions()), mj_model, model)))\n",
-    "    model.integrate(t0=0.0, tf=timestep, integrator_type=IntegratorType.EulerSemiImplicit)\n",
+    "    sim_images.append(\n",
+    "        get_image(\n",
+    "            \"side\",\n",
+    "            from_jaxsim_to_mujoco_pos(\n",
+    "                np.array(model.joint_positions()), mj_model, model\n",
+    "            ),\n",
+    "        )\n",
+    "    )\n",
+    "    model.integrate(\n",
+    "        t0=0.0, tf=timestep, integrator_type=IntegratorType.EulerSemiImplicit\n",
+    "    )\n",
     "\n",
-    "media.show_video(sim_images, fps=1/timestep)"
+    "media.show_video(sim_images, fps=1 / timestep)"
    ]
   },
   {
@@ -274,6 +293,7 @@
     "# Compute the gravity compensation term\n",
     "H = model.free_floating_bias_forces()[6:]\n",
     "\n",
+    "\n",
     "def pd_controller(\n",
     "    q: jax.Array, q_d: jax.Array, q_dot: jax.Array, q_dot_d: jax.Array\n",
     ") -> jax.Array:\n",
@@ -297,7 +317,14 @@
     "timestep = 0.01\n",
     "\n",
     "for _ in range(300):\n",
-    "    sim_images.append(get_image(\"side\", from_jaxsim_to_mujoco_pos(np.array(model.joint_positions()), mj_model, model)))\n",
+    "    sim_images.append(\n",
+    "        get_image(\n",
+    "            \"side\",\n",
+    "            from_jaxsim_to_mujoco_pos(\n",
+    "                np.array(model.joint_positions()), mj_model, model\n",
+    "            ),\n",
+    "        )\n",
+    "    )\n",
     "    model.set_joint_generalized_force_targets(\n",
     "        forces=pd_controller(\n",
     "            q=model.joint_positions(),\n",
@@ -306,9 +333,11 @@
     "            q_dot_d=jnp.array([0.0, 0.0]),\n",
     "        )\n",
     "    )\n",
-    "    model.integrate(t0=0.0, tf=timestep, integrator_type=IntegratorType.EulerSemiImplicit)\n",
+    "    model.integrate(\n",
+    "        t0=0.0, tf=timestep, integrator_type=IntegratorType.EulerSemiImplicit\n",
+    "    )\n",
     "\n",
-    "media.show_video(sim_images, fps=1/timestep)"
+    "media.show_video(sim_images, fps=1 / timestep)"
    ]
   }
  ],

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/examples/Parallel_computing.ipynb

@@ -24,7 +24,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#@title Imports and setup\n",
+    "# @title Imports and setup\n",
     "import sys\n",
     "\n",
     "from IPython.display import HTML, clear_output, display\n",
@@ -71,7 +71,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#@title Create a sphere model\n",
+    "# @title Create a sphere model\n",
     "model_sdf_string = rod.Sdf(\n",
     "    version=\"1.7\",\n",
     "    model=SphereBuilder(radius=0.10, mass=1.0, name=\"sphere\")\n",
@@ -119,9 +119,7 @@
     "\n",
     "\n",
     "# Add model to simulator\n",
-    "model = simulator.insert_model_from_description(\n",
-    "    model_description=model_sdf_string\n",
-    ")"
+    "model = simulator.insert_model_from_description(model_description=model_sdf_string)"
    ]
   },
   {

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/setup.cfg

@@ -62,7 +62,7 @@ where = src
 
 [options.extras_require]
 style = 
-	black
+	black[jupyter]
 	isort
 testing = 
 	idyntree

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.dev2'
-__version_tuple__ = version_tuple = (0, 2, 'dev2')
+__version__ = version = '0.2.dev8'
+__version_tuple__ = version_tuple = (0, 2, 'dev8')

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/src/jaxsim/high_level/model.py

@@ -385,13 +385,13 @@ class Model(Vmappable):
     def link_names(self) -> tuple[str, ...]:
         """"""
 
-        return tuple(l.name() for l in self.links())
+        return tuple(self.physics_model.description.links_dict.keys())
 
     @functools.partial(oop.jax_tf.method_ro, jit=False, vmap=False)
     def joint_names(self) -> tuple[str, ...]:
         """"""
 
-        return tuple(j.name() for j in self.joints())
+        return tuple(self.physics_model.description.joints_dict.keys())
 
     @functools.partial(oop.jax_tf.method_ro, jit=False, vmap=False)
     def links(

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/src/jaxsim/parsers/rod/parser.py

@@ -135,11 +135,13 @@ def extract_model_data(
                 parent=world_link,
                 child=links_dict[j.child],
                 jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
-                axis=np.array(j.axis.xyz.xyz)
-                if j.axis is not None
-                and j.axis.xyz is not None
-                and j.axis.xyz.xyz is not None
-                else None,
+                axis=(
+                    np.array(j.axis.xyz.xyz)
+                    if j.axis is not None
+                    and j.axis.xyz is not None
+                    and j.axis.xyz.xyz is not None
+                    else None
+                ),
                 pose=j.pose.transform() if j.pose is not None else np.eye(4),
             )
             for j in sdf_model.joints()
@@ -200,41 +202,55 @@ def extract_model_data(
             parent=links_dict[j.parent],
             child=links_dict[j.child],
             jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
-            axis=np.array(j.axis.xyz.xyz)
-            if j.axis is not None
-            and j.axis.xyz is not None
-            and j.axis.xyz.xyz is not None
-            else None,
+            axis=(
+                np.array(j.axis.xyz.xyz)
+                if j.axis is not None
+                and j.axis.xyz is not None
+                and j.axis.xyz.xyz is not None
+                else None
+            ),
             pose=j.pose.transform() if j.pose is not None else np.eye(4),
             initial_position=0.0,
             position_limit=(
-                float(j.axis.limit.lower)
-                if j.axis is not None and j.axis.limit is not None
-                else np.finfo(float).min,
-                float(j.axis.limit.upper)
-                if j.axis is not None and j.axis.limit is not None
-                else np.finfo(float).max,
+                (
+                    float(j.axis.limit.lower)
+                    if j.axis is not None and j.axis.limit is not None
+                    else np.finfo(float).min
+                ),
+                (
+                    float(j.axis.limit.upper)
+                    if j.axis is not None and j.axis.limit is not None
+                    else np.finfo(float).max
+                ),
+            ),
+            friction_static=(
+                j.axis.dynamics.friction
+                if j.axis is not None
+                and j.axis.dynamics is not None
+                and j.axis.dynamics.friction is not None
+                else 0.0
+            ),
+            friction_viscous=(
+                j.axis.dynamics.damping
+                if j.axis is not None
+                and j.axis.dynamics is not None
+                and j.axis.dynamics.damping is not None
+                else 0.0
+            ),
+            position_limit_damper=(
+                j.axis.limit.dissipation
+                if j.axis is not None
+                and j.axis.limit is not None
+                and j.axis.limit.dissipation is not None
+                else 0.0
+            ),
+            position_limit_spring=(
+                j.axis.limit.stiffness
+                if j.axis is not None
+                and j.axis.limit is not None
+                and j.axis.limit.stiffness is not None
+                else 0.0
             ),
-            friction_static=j.axis.dynamics.friction
-            if j.axis is not None
-            and j.axis.dynamics is not None
-            and j.axis.dynamics.friction is not None
-            else 0.0,
-            friction_viscous=j.axis.dynamics.damping
-            if j.axis is not None
-            and j.axis.dynamics is not None
-            and j.axis.dynamics.damping is not None
-            else 0.0,
-            position_limit_damper=j.axis.limit.dissipation
-            if j.axis is not None
-            and j.axis.limit is not None
-            and j.axis.limit.dissipation is not None
-            else 0.0,
-            position_limit_spring=j.axis.limit.stiffness
-            if j.axis is not None
-            and j.axis.limit is not None
-            and j.axis.limit.stiffness is not None
-            else 0.0,
         )
         for j in sdf_model.joints()
         if j.type in {"revolute", "prismatic", "fixed"}

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/src/jaxsim/physics/model/physics_model.py

@@ -45,14 +45,14 @@ class PhysicsModel(JaxsimDataclass):
     )
     is_floating_base: Static[bool] = dataclasses.field(default=False)
     gc: GroundContact = dataclasses.field(default_factory=lambda: GroundContact())
-    description: Static[
-        jaxsim.parsers.descriptions.model.ModelDescription
-    ] = dataclasses.field(default=None)
+    description: Static[jaxsim.parsers.descriptions.model.ModelDescription] = (
+        dataclasses.field(default=None)
+    )
 
     _parent_array_dict: Static[Dict[int, int]] = dataclasses.field(default_factory=dict)
-    _jtype_dict: Static[
-        Dict[int, Union[JointType, JointDescriptor]]
-    ] = dataclasses.field(default_factory=dict)
+    _jtype_dict: Static[Dict[int, Union[JointType, JointDescriptor]]] = (
+        dataclasses.field(default_factory=dict)
+    )
     _tree_transforms_dict: Dict[int, jtp.Matrix] = dataclasses.field(
         default_factory=dict
     )

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/src/jaxsim/simulation/simulator.py

@@ -432,8 +432,9 @@ class JaxSim(Vmappable):
     def step_over_horizon(
         self,
         horizon_steps: jtp.Int,
-        callback_handler: Union["scb.SimulatorCallback", "scb.CallbackHandler"]
-        | None = None,
+        callback_handler: (
+            Union["scb.SimulatorCallback", "scb.CallbackHandler"] | None
+        ) = None,
         clear_inputs: jtp.Bool = False,
     ) -> Union[
         "JaxSim",
@@ -459,10 +460,8 @@ class JaxSim(Vmappable):
         sim = self.copy().mutable(validate=True)
 
         # Helper to get callbacks from the handler
-        get_cb = (
-            lambda h, cb_name: getattr(h, cb_name)
-            if h is not None and hasattr(h, cb_name)
-            else None
+        get_cb = lambda h, cb_name: (
+            getattr(h, cb_name) if h is not None and hasattr(h, cb_name) else None
         )
 
         # Get the callbacks

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/src/jaxsim/utils/oop.py

@@ -3,16 +3,20 @@ import dataclasses
 import functools
 import inspect
 import os
-from typing import Any, Callable, Generator
+from typing import Any, Callable, Generator, TypeVar
 
 import jax
 import jax.flatten_util
+from typing_extensions import ParamSpec
 
 from jaxsim import logging
 from jaxsim.utils import tracing
 
 from . import Mutability, Vmappable
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
 
 class jax_tf:
     """
@@ -27,13 +31,13 @@ class jax_tf:
 
     @staticmethod
     def method_ro(
-        fn: Callable,
+        fn: Callable[_P, _R],
         jit: bool = True,
         static_argnames: tuple[str, ...] | list[str] = (),
         vmap: bool | None = None,
         vmap_in_axes: tuple[int, ...] | int | None = None,
         vmap_out_axes: tuple[int, ...] | int | None = None,
-    ):
+    ) -> Callable[_P, _R]:
         """
         Decorator for r/o methods of classes inheriting from Vmappable.
         """
@@ -51,14 +55,14 @@ class jax_tf:
 
     @staticmethod
     def method_rw(
-        fn: Callable,
+        fn: Callable[_P, _R],
         validate: bool = True,
         jit: bool = True,
         static_argnames: tuple[str, ...] | list[str] = (),
         vmap: bool | None = None,
         vmap_in_axes: tuple[int, ...] | int | None = None,
         vmap_out_axes: tuple[int, ...] | int | None = None,
-    ):
+    ) -> Callable[_P, _R]:
         """
         Decorator for r/w methods of classes inheriting from Vmappable.
         """
@@ -76,7 +80,7 @@ class jax_tf:
 
     @staticmethod
     def method(
-        fn: Callable,
+        fn: Callable[_P, _R],
         read_only: bool = True,
         validate: bool = True,
         jit_enabled: bool = True,
@@ -109,7 +113,7 @@ class jax_tf:
         """
 
         @functools.wraps(fn)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: _P.args, **kwargs: _P.kwargs):
             """The wrapper function that is returned by this decorator."""
 
             # Methods of classes inheriting from Vmappable decorated by this wrapper
@@ -202,9 +206,9 @@ class jax_tf:
             mutability_dict = {
                 Mutability.MUTABLE_NO_VALIDATION: Mutability.MUTABLE_NO_VALIDATION,
                 Mutability.MUTABLE: Mutability.MUTABLE,
-                Mutability.FROZEN: Mutability.MUTABLE
-                if validate
-                else Mutability.MUTABLE_NO_VALIDATION,
+                Mutability.FROZEN: (
+                    Mutability.MUTABLE if validate else Mutability.MUTABLE_NO_VALIDATION
+                ),
             }
 
             # We need to replace all the dynamic leafs of the original instance with those

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.dev2
+Version: 0.2.dev8
 Summary: A physics engine in reduced coordinates implemented with JAX.
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
@@ -38,7 +38,7 @@ Requires-Dist: jax_dataclasses>=1.4.0
 Requires-Dist: pptree
 Requires-Dist: rod
 Provides-Extra: style
-Requires-Dist: black; extra == "style"
+Requires-Dist: black[jupyter]; extra == "style"
 Requires-Dist: isort; extra == "style"
 Provides-Extra: testing
 Requires-Dist: idyntree; extra == "testing"
@@ -47,7 +47,7 @@ Requires-Dist: pytest-forked; extra == "testing"
 Requires-Dist: pytest-icdiff; extra == "testing"
 Requires-Dist: robot-descriptions; extra == "testing"
 Provides-Extra: all
-Requires-Dist: black; extra == "all"
+Requires-Dist: black[jupyter]; extra == "all"
 Requires-Dist: isort; extra == "all"
 Requires-Dist: idyntree; extra == "all"
 Requires-Dist: pytest>=6.0; extra == "all"

{jaxsim-0.2.dev2 → jaxsim-0.2.dev8}/src/jaxsim.egg-info/requires.txt

@@ -7,7 +7,7 @@ pptree
 rod
 
 [all]
-black
+black[jupyter]
 isort
 idyntree
 pytest>=6.0
@@ -16,7 +16,7 @@ pytest-icdiff
 robot-descriptions
 
 [style]
-black
+black[jupyter]
 isort
 
 [testing]