Merged comments from Junpeng

blackjax-devs · Sep 16, 2024 · 35442a3 · 35442a3
1 parent 1cbe9cf
commit 35442a3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/blackjax/mcmc/metrics.py b/blackjax/mcmc/metrics.py
@@ -187,7 +187,7 @@ def is_turning(
         return turning_at_left | turning_at_right
 
     def scale(
-        position: ArrayLikeTree, element: ArrayLikeTree, inv: bool
+        position: ArrayLikeTree, element: ArrayLikeTree, inv: ArrayLikeTree
     ) -> ArrayLikeTree:
         """Scale elements by the mass matrix.
 
@@ -279,7 +279,7 @@ def is_turning(
         # return turning_at_left | turning_at_right
 
     def scale(
-        position: ArrayLikeTree, element: ArrayLikeTree, inv: bool
+        position: ArrayLikeTree, element: ArrayLikeTree, inv: ArrayLikeTree
     ) -> ArrayLikeTree:
         """Scale elements by the mass matrix.
 

diff --git a/blackjax/types.py b/blackjax/types.py
@@ -46,4 +46,4 @@ class WelfordAlgorithmState(NamedTuple):
 
 #: JAX Scalar types
 Scalar = Union[float, int]
-Numeric = Union[Array, Scalar]
+Numeric = Union[jax.Array, Scalar]