Update util.py

Updating types.
jax-md · Apr 8, 2024 · 3b8c58e · 3b8c58e
1 parent f1af70a
commit 3b8c58e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax_md/util.py b/jax_md/util.py
@@ -25,7 +25,7 @@
 
 import numpy as onp
 
-Array = jnp.ndarray
+Array = Any
 PyTree = Any
 
 i16 = jnp.int16