Fix MALA transition energy (#653)

* Fix MALA transition energy * Use a different logic.
blackjax-devs · Mar 31, 2024 · f77297f · f77297f
1 parent 2e25624
commit f77297f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/blackjax/mcmc/mala.py b/blackjax/mcmc/mala.py
@@ -80,14 +80,14 @@ def transition_energy(state, new_state, step_size):
         """Transition energy to go from `state` to `new_state`"""
         theta = jax.tree_util.tree_map(
             lambda new_x, x, g: new_x - x - step_size * g,
-            new_state.position,
             state.position,
-            state.logdensity_grad,
+            new_state.position,
+            new_state.logdensity_grad,
         )
         theta_dot = jax.tree_util.tree_reduce(
             operator.add, jax.tree_util.tree_map(lambda x: jnp.sum(x * x), theta)
         )
-        return -state.logdensity + 0.25 * (1.0 / step_size) * theta_dot
+        return -new_state.logdensity + 0.25 * (1.0 / step_size) * theta_dot
 
     compute_acceptance_ratio = proposal.compute_asymmetric_acceptance_ratio(
         transition_energy