From c8d6b534c2f0524f0ef6c1207a4ff4f00ac9b229 Mon Sep 17 00:00:00 2001
From: Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
Date: Tue, 22 Oct 2024 18:45:36 +0200
Subject: [PATCH] Use bitwise operators instead of `jnp.logical_*`

---
 src/jaxsim/api/model.py                 | 18 +++++++++---------
 src/jaxsim/integrators/variable_step.py | 12 ++++--------
 2 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/src/jaxsim/api/model.py b/src/jaxsim/api/model.py
index 0c9e63f62..7dad3606d 100644
--- a/src/jaxsim/api/model.py
+++ b/src/jaxsim/api/model.py
@@ -2027,15 +2027,15 @@ def step(
             # Raise runtime error for not supported case in which Rigid contacts and
             # Baumgarte stabilization are enabled and used with ForwardEuler integrator.
             jaxsim.exceptions.raise_runtime_error_if(
-                condition=jnp.logical_and(
-                    isinstance(
-                        integrator,
-                        jaxsim.integrators.fixed_step.ForwardEuler
-                        | jaxsim.integrators.fixed_step.ForwardEulerSO3,
-                    ),
-                    jnp.array(
-                        [data_tf.contacts_params.K, data_tf.contacts_params.D]
-                    ).any(),
+                condition=isinstance(
+                    integrator,
+                    jaxsim.integrators.fixed_step.ForwardEuler
+                    | jaxsim.integrators.fixed_step.ForwardEulerSO3,
+                )
+                & (
+                    data_tf.contacts_params.K
+                    is not None | data_tf.contacts_params.D
+                    is not None
                 ),
                 msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
             )
diff --git a/src/jaxsim/integrators/variable_step.py b/src/jaxsim/integrators/variable_step.py
index 79c7fab06..875ed5bfe 100644
--- a/src/jaxsim/integrators/variable_step.py
+++ b/src/jaxsim/integrators/variable_step.py
@@ -446,14 +446,10 @@ def reject_step():
                 params_next,
                 discarded_steps,
             ) = jax.lax.cond(
-                pred=jnp.array(
-                    [
-                        discarded_steps >= self.max_step_rejections,
-                        local_error <= 1.0,
-                        Δt_next < self.dt_min,
-                        integrator_init,
-                    ]
-                ).any(),
+                pred=discarded_steps
+                >= self.max_step_rejections | local_error
+                <= 1.0 | Δt_next
+                < self.dt_min | integrator_init,
                 true_fun=accept_step,
                 false_fun=reject_step,
             )