Update PyTree type definition

ami-iit · Mar 6, 2024 · 53bac0b · 53bac0b
1 parent 3b0077c
commit 53bac0b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/jaxsim/typing.py b/src/jaxsim/typing.py
@@ -15,7 +15,9 @@
 VectorJax = ArrayJax
 MatrixJax = ArrayJax
 
-PyTree = dict[Hashable, "PyTree"] | list["PyTree"] | tuple["PyTree"] | None | Any
+PyTree = (
+    dict[Hashable, "PyTree"] | list["PyTree"] | tuple["PyTree"] | None | jax.Array | Any
+)
 
 # =======================
 # Mixed JAX / NumPy types