Update typing module

src/jaxsim/typing.py

@@ -1,39 +1,37 @@
-from typing import Any, Dict, Hashable, List, NamedTuple, Tuple, Union
+from typing import Any, Hashable
 
-import jax.numpy as jnp
-import numpy.typing as npt
+import jax
 
+# =========
 # JAX types
-FloatJax = Union[jnp.float16, jnp.float32, jnp.float64]
-IntJax = Union[
-    jnp.int8,
-    jnp.int16,
-    jnp.int32,
-    jnp.int64,
-    jnp.uint8,
-    jnp.uint16,
-    jnp.uint32,
-    jnp.uint64,
-]
-ArrayJax = jnp.ndarray
-TensorJax = jnp.ndarray
+# =========
+
+ScalarJax = jax.Array
+IntJax = ScalarJax
+BoolJax = ScalarJax
+FloatJax = ScalarJax
+
+ArrayJax = jax.Array
 VectorJax = ArrayJax
 MatrixJax = ArrayJax
-PyTree = Union[
-    TensorJax,
-    Dict[Hashable, "PyTree"],
-    List["PyTree"],
-    NamedTuple,
-    Tuple["PyTree"],
-    None,
-    Any,
-]
 
+PyTree = dict[Hashable, "PyTree"] | list["PyTree"] | tuple["PyTree"] | None | Any
+
+# =======================
 # Mixed JAX / NumPy types
-Array = Union[npt.NDArray, ArrayJax]
-Tensor = Union[npt.NDArray, ArrayJax]
+# =======================
+
+Array = jax.typing.ArrayLike
 Vector = Array
 Matrix = Array
-Bool = Union[bool, ArrayJax]
-Int = Union[int, IntJax]
-Float = Union[float, FloatJax]
+
+Int = int | IntJax
+Bool = bool | ArrayJax
+Float = float | FloatJax
+
+ArrayLike = Array
+VectorLike = Vector
+MatrixLike = Matrix
+IntLike = Int
+BoolLike = Bool
+FloatLike = Float