Merge branch 'main' into kernel_refactoring

blackjax-devs · Mar 1, 2023 · 5c2a920 · 5c2a920
2 parents cf45d3b + 1908453
commit 5c2a920
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/blackjax/types.py b/blackjax/types.py
@@ -1,11 +1,7 @@
 from typing import Any, Iterable, Mapping, Union
 
 import jax
-import jax.numpy as jnp
-import numpy as np
-
-#: JAX or Numpy array
-Array = Union[np.ndarray, jnp.ndarray]
+from chex import Array
 
 #: JAX PyTrees
 PyTree = Union[Array, Iterable["PyTree"], Mapping[Any, "PyTree"]]