SPLIT TUNING FOR AMCLMC INTO SEPARATE FILE

blackjax-devs · Jun 4, 2024 · 31a7be6 · 31a7be6
1 parent da83d56
commit 31a7be6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/blackjax/util.py b/blackjax/util.py
@@ -220,7 +220,7 @@ def one_step(average_and_state, xs, return_state):
         if return_state:
             return (average, state), (transform(state), info)
         else:
-            return (average, state), average[1]
+            return (average, state), None
 
     one_step = jax.jit(partial(one_step, return_state=return_state_history))
 
@@ -233,7 +233,7 @@ def one_step(average_and_state, xs, return_state):
     )
 
     if not return_state_history:
-        return average, transform(final_state), history
+        return average, transform(final_state)
     else:
         state_history, info_history = history
         return transform(final_state), state_history, info_history