Merge branch 'emaus' of github.com:reubenharry/blackjax into emaus

Author: reubenharry

blackjax/adaptation/ensemble_mclmc.py

@@ -71,7 +71,7 @@ def __init__(
         acc_prob_target=0.8,
         observables=lambda x: 0.0,  # just for diagnostics: some function of a given chain at given timestep
         observables_for_bias=lambda x: 0.0,  # just for diagnostics: the above, but averaged over all chains
-        contract=lambda x: 0.0,  # just for diagnostics: observabiels for bias, contracted over dimensions
+        contract=lambda x: 0.0,  # just for diagnostics: observables for bias, contracted over dimensions
     ):
         self.num_adaptation_samples = num_adaptation_samples
         self.observables = observables
@@ -111,7 +111,7 @@ def summary_statistics_fn(self, state, info, rng_key):
     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"])  # remove
+        true_bias = self.contract(Etheta["observables_for_bias"]) 
 
         info_to_be_stored = {
             "L": adaptation_state.step_size * adaptation_state.steps_per_sample,
@@ -164,7 +164,6 @@ def while_steps_num(cond):
 def emaus(
     logdensity_fn,
     sample_init,
-    transform,
     ndims,
     num_steps1,
     num_steps2,
@@ -180,9 +179,10 @@ def emaus(
     integrator_coefficients=None,
     steps_per_sample=15,
     acc_prob=None,
-    observables=lambda x: None,
+    observables_for_bias=lambda x: 0.0,
     ensemble_observables=None,
     diagnostics=True,
+    contract=lambda x: 0.0,
 ):
     """
     model: the target density object
@@ -215,7 +215,7 @@ def emaus(
 
     # burn-in with the unadjusted method #
     kernel = umclmc.build_kernel(logdensity_fn)
-    save_num = (int)(jnp.rint(save_frac * num_steps1))
+    save_num = (jnp.rint(save_frac * num_steps1)).astype(int)
     adap = umclmc.Adaptation(
         ndims,
         alpha=alpha,
@@ -224,9 +224,8 @@ def emaus(
         C=C,
         power=3.0 / 8.0,
         r_end=r_end,
-        observables_for_bias=lambda position: jnp.square(
-            transform(jax.flatten_util.ravel_pytree(position)[0])
-        ),
+        observables_for_bias=observables_for_bias,
+        contract=contract,
     )
 
     final_state, final_adaptation_state, info1 = run_eca(
@@ -294,8 +293,11 @@ def emaus(
         num_adaptation_samples,
         steps_per_sample,
         _acc_prob,
+        contract=contract,
+        observables_for_bias=observables_for_bias,
     )
 
+
     final_state, final_adaptation_state, info2 = run_eca(
         key_mclmc,
         initial_state,

blackjax/adaptation/ensemble_umclmc.py

@@ -72,6 +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(
@@ -82,9 +83,8 @@ def sequential_init(key, x, args):
 
     def summary_statistics_fn(state):
         """compute the diagonal elements of the equipartition matrix"""
-        return 0  # -state.position * state.logdensity_grad
+        return -state.position * state.logdensity_grad
 
-    # TODO: restore!
 
     def ensemble_init(key, state, signs):
         """flip the velocity, depending on the equipartition condition"""

blackjax/adaptation/mclmc_adaptation.py

@@ -50,8 +50,8 @@ def mclmc_find_L_and_step_size(
     desired_energy_var=5e-4,
     trust_in_estimate=1.5,
     num_effective_samples=150,
-    diagonal_preconditioning=True,
     params=None,
+    diagonal_preconditioning=True,
 ):
     """
     Finds the optimal value of the parameters for the MCLMC algorithm.

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,11 +352,14 @@ 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), (
@@ -381,6 +384,7 @@ def run_eca(
     ensemble_info=None,
     early_stop=False,
 ):
+    
     """
     Run ensemble chain adaptation (eca) in parallel on multiple devices.
     -----------------------------------------------------
@@ -413,6 +417,7 @@ def all_steps(initial_state, keys_sampling, keys_adaptation):
 
         initial_state_all = (initial_state, adaptation.initial_state)
 
+
         # run sampling
         xs = (
             jnp.arange(num_steps),
@@ -441,13 +446,16 @@ 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,