names fix

blackjax-devs · Oct 30, 2023 · 8b540bc · 8b540bc
1 parent ba86244
commit 8b540bc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/blackjax/adaptation/chees_adaptation.py b/blackjax/adaptation/chees_adaptation.py
@@ -74,7 +74,7 @@ def base(halton_sequence: Array, optim: optax.GradientTransformation):
         Function that initializes the warmup.
     update
         Function that moves the warmup one step.
-    
+
     """
 
     da_init, da_update, _ = dual_averaging.dual_averaging()
@@ -279,7 +279,7 @@ def chees_adaptation(
 
     .. code::
 
-        warmup = blackjax.chees(logprob_fn, num_chains)
+        warmup = blackjax.chees_adaptation(logprob_fn, num_chains)
         key_warmup, key_sample = jax.random.split(rng_key)
         optim = optax.adam(learning_rate)
         last_states, kernel, parameters = warmup.run(

diff --git a/tests/mcmc/test_sampling.py b/tests/mcmc/test_sampling.py
@@ -238,15 +238,15 @@ def test_chees(self):
         init_key, warmup_key, inference_key = jax.random.split(rng_key, 3)
 
         num_chains = 128
-        warmup = blackjax.chees(
+        warmup = blackjax.chees_adaptation(
             logposterior_fn,
             num_chains=num_chains,
         )
         scale_key, coefs_key = jax.random.split(init_key, 2)
         log_scales = 1.0 + jax.random.normal(scale_key, (num_chains,))
         coefs = 4.0 + jax.random.normal(coefs_key, (num_chains,))
         initial_positions = {"log_scale": log_scales, "coefs": coefs}
-        last_states, kernel, _ = warmup.run(
+        (last_states, _), kernel, _ = warmup.run(
             warmup_key,
             initial_positions,
             step_size=0.001,