fix operations with pytrees

Harry45 · Jan 14, 2021 · facc554 · facc554
1 parent 30a82cc
commit facc554
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/blackjax/hmc.py b/blackjax/hmc.py
@@ -80,7 +80,7 @@ def kernel(potential_fn: Callable, parameters: HMCParameters) -> Callable:
     """
     step_size, num_integration_steps, inv_mass_matrix, divergence_threshold = parameters
 
-    if not inv_mass_matrix:
+    if inv_mass_matrix is None:
         raise ValueError(
             "Expected a value for `inv_mass_matrix`,"
             " got None. Please specify a value when initializing"

diff --git a/blackjax/inference/base.py b/blackjax/inference/base.py
@@ -152,7 +152,7 @@ def kernel(
         )
         new_position, new_momentum, new_potential_energy_grad = proposal
 
-        flipped_momentum = -1.0 * new_momentum
+        flipped_momentum = jax.tree_util.tree_multimap(lambda m: -1.0 * m, new_momentum)
         new_potential_energy = potential_fn(new_position)
         new_energy = new_potential_energy + kinetic_energy(
             flipped_momentum, new_position

diff --git a/blackjax/inference/metrics.py b/blackjax/inference/metrics.py
@@ -68,6 +68,7 @@ def momentum_generator(rng_key: jax.random.PRNGKey, position: PyTree) -> PyTree:
 
         def kinetic_energy(momentum: PyTree, *_) -> float:
             momentum, _ = tree_flatten(momentum)
+            momentum = jnp.array(momentum)
             velocity = jnp.multiply(inverse_mass_matrix, momentum)
             return 0.5 * jnp.dot(velocity, momentum)
 
@@ -85,6 +86,7 @@ def momentum_generator(rng_key: jax.random.PRNGKey, position: PyTree) -> PyTree:
 
         def kinetic_energy(momentum: PyTree, *_) -> float:
             momentum, _ = tree_flatten(momentum)
+            momentum = jnp.array(momentum)
             velocity = jnp.matmul(inverse_mass_matrix, momentum)
             return 0.5 * jnp.dot(velocity, momentum)