Update blackjax/adaptation/adjusted_mclmc_adaptation.py

Co-authored-by: Junpeng Lao <[email protected]>
blackjax-devs · Jan 1, 2025 · a75589b · a75589b
1 parent aa98579
commit a75589b
Showing 1 changed file with 4 additions and 7 deletions.
diff --git a/blackjax/adaptation/adjusted_mclmc_adaptation.py b/blackjax/adaptation/adjusted_mclmc_adaptation.py
@@ -192,14 +192,11 @@ def step(iteration_state, weight_and_key):
                 zero_prevention=mask,
             )
 
-            if fix_L:
-                params = params._replace(
-                    step_size=mask * step_size + (1 - mask) * params.step_size,
-                )
-
-            else:
+            params = params._replace(
+                step_size=mask * step_size + (1 - mask) * params.step_size,
+            )
+            if not fix_L:
                 params = params._replace(
-                    step_size=mask * step_size + (1 - mask) * params.step_size,
                     L=mask * (params.L * (step_size / params.step_size))
                     + (1 - mask) * params.L,
                 )