From 50aea065d4ed3825b8022b7d622d7f7762270e2a Mon Sep 17 00:00:00 2001
From: diegoferigo <diego.ferigo@iit.it>
Date: Wed, 6 Dec 2023 05:18:25 +0100
Subject: [PATCH] Update bool typing

---
 src/jaxsim/typing.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/jaxsim/typing.py b/src/jaxsim/typing.py
index 82355cf33..94b9508be 100644
--- a/src/jaxsim/typing.py
+++ b/src/jaxsim/typing.py
@@ -1,7 +1,6 @@
 from typing import Any, Dict, Hashable, List, NamedTuple, Tuple, Union
 
 import jax.numpy as jnp
-import numpy as np
 import numpy.typing as npt
 
 # JAX types
@@ -35,6 +34,6 @@
 Tensor = Union[npt.NDArray, ArrayJax]
 Vector = Array
 Matrix = Array
-Bool = bool
+Bool = Union[bool, ArrayJax]
 Int = Union[int, IntJax]
 Float = Union[float, FloatJax]