Update bool typing

flferretti · Dec 6, 2023 · 50aea06 · 50aea06
1 parent 458eeac
commit 50aea06
Showing 1 changed file with 1 addition and 2 deletions.
diff --git 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]