From 7cf4f9d621917d7bafa341653a6776cca66c29fb Mon Sep 17 00:00:00 2001
From: ksnxr <70186663+ksnxr@users.noreply.github.com>
Date: Mon, 1 Apr 2024 09:46:45 +0300
Subject: [PATCH] Change variable names (#654)

---
 blackjax/mcmc/mala.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/blackjax/mcmc/mala.py b/blackjax/mcmc/mala.py
index f6dd7c106..1f1345cc4 100644
--- a/blackjax/mcmc/mala.py
+++ b/blackjax/mcmc/mala.py
@@ -79,7 +79,7 @@ def build_kernel():
     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,
+            lambda x, new_x, g: x - new_x - step_size * g,
             state.position,
             new_state.position,
             new_state.logdensity_grad,