bug fix

Author: reubenharry

blackjax/adaptation/ensemble_mclmc.py

@@ -101,17 +101,15 @@ def __init__(
     def summary_statistics_fn(self, state, info, rng_key):
         return {
             "acceptance_probability": info.acceptance_rate,
-            "equipartition_diagonal": equipartition_diagonal(
-                state
-            ),  # metric for bias: equipartition theorem gives todo...
+            "equipartition_diagonal": equipartition_diagonal(state),
             "observables": self.observables(state.position),
             "observables_for_bias": self.observables_for_bias(state.position),
         }
 
     def update(self, adaptation_state, Etheta):
         acc_prob = Etheta["acceptance_probability"]
         equi_diag = equipartition_diagonal_loss(Etheta["equipartition_diagonal"])
-        true_bias = self.contract(Etheta["observables_for_bias"]) 
+        true_bias = self.contract(Etheta["observables_for_bias"])
 
         info_to_be_stored = {
             "L": adaptation_state.step_size * adaptation_state.steps_per_sample,
@@ -179,7 +177,7 @@ def emaus(
     integrator_coefficients=None,
     steps_per_sample=15,
     acc_prob=None,
-    observables_for_bias=lambda x: 0.0,
+    observables_for_bias=lambda x: x,
     ensemble_observables=None,
     diagnostics=True,
     contract=lambda x: 0.0,
@@ -205,7 +203,6 @@ def emaus(
     diagnostics: whether to return diagnostics
     """
 
-    # observables_for_bias, contract = bias(model)
     key_init, key_umclmc, key_mclmc = jax.random.split(rng_key, 3)
 
     # initialize the chains
@@ -297,7 +294,6 @@ def emaus(
         observables_for_bias=observables_for_bias,
     )
 
-
     final_state, final_adaptation_state, info2 = run_eca(
         key_mclmc,
         initial_state,

blackjax/adaptation/ensemble_umclmc.py

@@ -72,7 +72,7 @@ def initialize(rng_key, logdensity_fn, sample_init, num_chains, mesh):
     def sequential_init(key, x, args):
         """initialize the position using sample_init and the velocity along the gradient"""
         position = sample_init(key)
-        
+
         logdensity, logdensity_grad = jax.value_and_grad(logdensity_fn)(position)
         flat_g, unravel_fn = ravel_pytree(logdensity_grad)
         velocity = unravel_fn(
@@ -83,8 +83,12 @@ def sequential_init(key, x, args):
 
     def summary_statistics_fn(state):
         """compute the diagonal elements of the equipartition matrix"""
-        return -state.position * state.logdensity_grad
+        flat_pos, unflatten = jax.flatten_util.ravel_pytree(state.position)
+        flat_g, unravel_fn = ravel_pytree(state.logdensity_grad)
+        return unravel_fn(-flat_pos * flat_g)
+        # return 0
 
+    # -state.position # * state.logdensity_grad
 
     def ensemble_init(key, state, signs):
         """flip the velocity, depending on the equipartition condition"""
@@ -113,7 +117,9 @@ def ensemble_init(key, state, signs):
         summary_statistics_fn=summary_statistics_fn,
     )
 
-    signs = -2.0 * (equipartition < 1.0) + 1.0
+    flat_equi, _ = ravel_pytree(equipartition)
+
+    signs = -2.0 * (flat_equi < 1.0) + 1.0
     initial_state, _ = ensemble_execute_fn(
         ensemble_init, key2, num_chains, mesh, x=initial_state, args=signs
     )
@@ -122,7 +128,6 @@ def ensemble_init(key, state, signs):
 
 
 def update_history(new_vals, history):
-    new_vals, _ = jax.flatten_util.ravel_pytree(new_vals)
     return jnp.concatenate((new_vals[None, :], history[:-1, :]))
 
 
@@ -258,7 +263,6 @@ def update(self, adaptation_state, Etheta):
         history_observables = update_history(
             Etheta["observables_for_bias"], adaptation_state.history.observables
         )
-        # history_observables = adaptation_state.history.observables
 
         history_weights = update_history_scalar(1.0, adaptation_state.history.weights)
         fluctuations = contract_history(history_observables, history_weights)

blackjax/util.py

@@ -10,7 +10,7 @@
 from jax.random import normal, split
 from jax.sharding import NamedSharding, PartitionSpec
 from jax.tree_util import tree_leaves, tree_map
-import jax
+
 from blackjax.base import SamplingAlgorithm, VIAlgorithm
 from blackjax.progress_bar import gen_scan_fn
 from blackjax.types import Array, ArrayLikeTree, ArrayTree, PRNGKey
@@ -352,14 +352,11 @@ def _step(state_all, xs):
         adaptation_state, info_to_be_stored = adaptation_update(
             adaptation_state, Etheta
         )
-        
 
         return (state, adaptation_state), info_to_be_stored
-    
 
     if ensemble_info is not None:
 
-
         def step(state_all, xs):
             (state, adaptation_state), info_to_be_stored = _step(state_all, xs)
             return (state, adaptation_state), (
@@ -384,7 +381,6 @@ def run_eca(
     ensemble_info=None,
     early_stop=False,
 ):
-    
     """
     Run ensemble chain adaptation (eca) in parallel on multiple devices.
     -----------------------------------------------------
@@ -417,7 +413,6 @@ def all_steps(initial_state, keys_sampling, keys_adaptation):
 
         initial_state_all = (initial_state, adaptation.initial_state)
 
-
         # run sampling
         xs = (
             jnp.arange(num_steps),
@@ -446,16 +441,13 @@ def step_while(a):
         else:
             final_state_all, info_history = lax.scan(step, initial_state_all, xs)
 
-        
-
         final_state, final_adaptation_state = final_state_all
         return (
             final_state,
             final_adaptation_state,
             info_history,
         )  # info history is composed of averages over all chains, so it is a couple of scalars
 
-
     p, pscalar = PartitionSpec("chains"), PartitionSpec()
     parallel_execute = shard_map(
         all_steps,

tests/mcmc/test_sampling.py

@@ -5,6 +5,7 @@
 
 import chex
 import jax
+
 # jax.config.update("jax_traceback_filtering", "off")
 import jax.numpy as jnp
 import jax.scipy.stats as stats
@@ -296,11 +297,10 @@ def run_emaus(
         sample_init,
         logdensity_fn,
         ndims,
-        transform,
         key,
         diagonal_preconditioning,
     ):
-        mesh = jax.sharding.Mesh(devices=jax.devices(),axis_names= "chains")
+        mesh = jax.sharding.Mesh(devices=jax.devices(), axis_names="chains")
 
         from blackjax.mcmc.integrators import mclachlan_coefficients
 
@@ -309,7 +309,6 @@ def run_emaus(
         info, grads_per_step, _acc_prob, final_state = emaus(
             logdensity_fn=logdensity_fn,
             sample_init=sample_init,
-            transform=transform,
             ndims=ndims,
             num_steps1=100,
             num_steps2=300,
@@ -602,7 +601,6 @@ def sample_init(key):
         samples = self.run_emaus(
             sample_init=sample_init,
             logdensity_fn=logdensity_fn,
-            transform=lambda x: x,
             ndims=2,
             key=inference_key,
             diagonal_preconditioning=True,