Simplify index extraction checks in api.link and api.joint

ami-iit · Oct 22, 2024 · 4742ae4 · 4742ae4
1 parent f3498cd
commit 4742ae4
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 8 deletions.
diff --git a/src/jaxsim/api/joint.py b/src/jaxsim/api/joint.py
@@ -53,9 +53,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, joint_index: jtp.IntLike) -> str
     """
 
     exceptions.raise_value_error_if(
-        condition=jnp.array(
-            [joint_index < 0, joint_index >= model.number_of_joints()]
-        ).any(),
+        condition=joint_index < 0,
         msg="Invalid joint index '{idx}'",
         idx=joint_index,
     )

diff --git a/src/jaxsim/api/link.py b/src/jaxsim/api/link.py
@@ -54,9 +54,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, link_index: jtp.IntLike) -> str:
     """
 
     exceptions.raise_value_error_if(
-        condition=jnp.array(
-            [link_index < 0, link_index >= model.number_of_links()]
-        ).any(),
+        condition=link_index < 0,
         msg="Invalid link index '{idx}'",
         idx=link_index,
     )

diff --git a/tests/test_api_joint.py b/tests/test_api_joint.py
@@ -41,5 +41,5 @@ def test_joint_index(
     with pytest.raises(jaxlib.xla_extension.XlaRuntimeError):
         _ = js.joint.idx_to_name(model=model, joint_index=-1)
 
-    with pytest.raises(jaxlib.xla_extension.XlaRuntimeError):
+    with pytest.raises(IndexError):
         _ = js.joint.idx_to_name(model=model, joint_index=model.number_of_joints())
diff --git a/tests/test_api_link.py b/tests/test_api_link.py
@@ -46,7 +46,7 @@ def test_link_index(
     with pytest.raises(jaxlib.xla_extension.XlaRuntimeError):
         _ = js.link.idx_to_name(model=model, link_index=-1)
 
-    with pytest.raises(jaxlib.xla_extension.XlaRuntimeError):
+    with pytest.raises(IndexError):
         _ = js.link.idx_to_name(model=model, link_index=model.number_of_links())