def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs):
rng_key, rng_r = random.split(rng_key)
state = super().init(rng_key, num_warmup, init_params, model_args,
model_kwargs)
self._support_sizes_flat, _ = ravel_pytree(
{k: self._support_sizes[k]
for k in self._gibbs_sites})
if self._num_discrete_updates is None:
self._num_discrete_updates = self._support_sizes_flat.shape[0]
self._num_warmup = num_warmup
# NB: the warmup adaptation can not be performed in sub-trajectories (i.e. the hmc trajectory
# between two discrete updates), so we will do it here, at the end of each MixedHMC step.
_, self._wa_update = warmup_adapter(
num_warmup,
adapt_step_size=self.inner_kernel._adapt_step_size,
adapt_mass_matrix=self.inner_kernel._adapt_mass_matrix,
dense_mass=self.inner_kernel._dense_mass,
target_accept_prob=self.inner_kernel._target_accept_prob,
find_reasonable_step_size=None,
)
# In HMC, when `hmc_state.r` is not None, we will skip drawing a random momemtum at the
# beginning of an HMC step. The reason is we need to maintain `r` between each sub-trajectories.
r = momentum_generator(state.hmc_state.z,
state.hmc_state.adapt_state.mass_matrix_sqrt,
rng_r)
return MixedHMCState(state.z, state.hmc_state._replace(r=r),
state.rng_key, jnp.zeros(()))
def sample(self, state, model_args, model_kwargs):
model_kwargs = {} if model_kwargs is None else model_kwargs
num_discretes = self._support_sizes_flat.shape[0]
def potential_fn(z_gibbs, z_hmc):
return self.inner_kernel._potential_fn_gen(*model_args,
_gibbs_sites=z_gibbs,
**model_kwargs)(z_hmc)
def update_discrete(idx, rng_key, hmc_state, z_discrete, ke_discrete,
delta_pe_sum):
# Algo 1, line 19: get a new discrete proposal
(
rng_key,
z_discrete_new,
pe_new,
log_accept_ratio,
) = self._discrete_proposal_fn(
rng_key,
z_discrete,
hmc_state.potential_energy,
partial(potential_fn, z_hmc=hmc_state.z),
idx,
self._support_sizes_flat[idx],
)
# Algo 1, line 20: depending on reject or refract, we will update
# the discrete variable and its corresponding kinetic energy. In case of
# refract, we will need to update the potential energy and its grad w.r.t. hmc_state.z
ke_discrete_i_new = ke_discrete[idx] + log_accept_ratio
grad_ = jacfwd if self.inner_kernel._forward_mode_differentiation else grad
z_discrete, pe, ke_discrete_i, z_grad = lax.cond(
ke_discrete_i_new > 0,
(z_discrete_new, pe_new, ke_discrete_i_new),
lambda vals: vals + (grad_(partial(potential_fn, vals[0]))
(hmc_state.z), ),
(
z_discrete,
hmc_state.potential_energy,
ke_discrete[idx],
hmc_state.z_grad,
),
identity,
)
delta_pe_sum = delta_pe_sum + pe - hmc_state.potential_energy
ke_discrete = ops.index_update(ke_discrete, idx, ke_discrete_i)
hmc_state = hmc_state._replace(potential_energy=pe, z_grad=z_grad)
return rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum
def update_continuous(hmc_state, z_discrete):
model_kwargs_ = model_kwargs.copy()
model_kwargs_["_gibbs_sites"] = z_discrete
hmc_state_new = self.inner_kernel.sample(hmc_state, model_args,
model_kwargs_)
# each time a sub-trajectory is performed, we need to reset i and adapt_state
# (we will only update them at the end of HMCGibbs step)
# For `num_steps`, we will record its cumulative sum for diagnostics
hmc_state = hmc_state_new._replace(
i=hmc_state.i,
adapt_state=hmc_state.adapt_state,
num_steps=hmc_state.num_steps + hmc_state_new.num_steps,
)
return hmc_state
def body_fn(i, vals):
(
rng_key,
hmc_state,
z_discrete,
ke_discrete,
delta_pe_sum,
arrival_times,
) = vals
idx = jnp.argmin(arrival_times)
# NB: length of each sub-trajectory is scaled from the current min(arrival_times)
# (see the note at total_time below)
trajectory_length = arrival_times[idx] * time_unit
arrival_times = arrival_times - arrival_times[idx]
arrival_times = ops.index_update(arrival_times, idx, 1.0)
# this is a trick, so that in a sub-trajectory of HMC, we always accept the new proposal
pe = jnp.inf
hmc_state = hmc_state._replace(trajectory_length=trajectory_length,
potential_energy=pe)
# Algo 1, line 7: perform a sub-trajectory
hmc_state = update_continuous(hmc_state, z_discrete)
# Algo 1, line 8: perform a discrete update
rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum = update_discrete(
idx, rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum)
return (
rng_key,
hmc_state,
z_discrete,
ke_discrete,
delta_pe_sum,
arrival_times,
)
z_discrete = {
k: v
for k, v in state.z.items() if k not in state.hmc_state.z
}
rng_key, rng_ke, rng_time, rng_r, rng_accept = random.split(
state.rng_key, 5)
# Algo 1, line 2: sample discrete kinetic energy
ke_discrete = random.exponential(rng_ke, (num_discretes, ))
# Algo 1, line 4 and 5: sample the initial amount of time that each discrete site visits
# the point 0/1. The logic in GetStepSizesNSteps(...) is more complicated but does
# the same job: the sub-trajectory length eta_t * M_t is the lag between two arrival time.
arrival_times = random.uniform(rng_time, (num_discretes, ))
# compute the amount of time to make `num_discrete_updates` discrete updates
total_time = (self._num_discrete_updates -
1) // num_discretes + jnp.sort(arrival_times)[
(self._num_discrete_updates - 1) % num_discretes]
# NB: total_time can be different from the HMC trajectory length, so we need to scale
# the time unit so that total_time * time_unit = hmc_trajectory_length
time_unit = state.hmc_state.trajectory_length / total_time
# Algo 1, line 2: sample hmc momentum
r = momentum_generator(state.hmc_state.r,
state.hmc_state.adapt_state.mass_matrix_sqrt,
rng_r)
hmc_state = state.hmc_state._replace(r=r, num_steps=0)
hmc_ke = euclidean_kinetic_energy(
hmc_state.adapt_state.inverse_mass_matrix, r)
# Algo 1, line 10: compute the initial energy
energy_old = hmc_ke + hmc_state.potential_energy
# Algo 1, line 3: set initial values
delta_pe_sum = 0.0
init_val = (
rng_key,
hmc_state,
z_discrete,
ke_discrete,
delta_pe_sum,
arrival_times,
)
# Algo 1, line 6-9: perform the update loop
rng_key, hmc_state_new, z_discrete_new, _, delta_pe_sum, _ = fori_loop(
0, self._num_discrete_updates, body_fn, init_val)
# Algo 1, line 10: compute the proposal energy
hmc_ke = euclidean_kinetic_energy(
hmc_state.adapt_state.inverse_mass_matrix, hmc_state_new.r)
energy_new = hmc_ke + hmc_state_new.potential_energy
# Algo 1, line 11: perform MH correction
delta_energy = energy_new - energy_old - delta_pe_sum
delta_energy = jnp.where(jnp.isnan(delta_energy), jnp.inf,
delta_energy)
accept_prob = jnp.clip(jnp.exp(-delta_energy), a_max=1.0)
# record the correct new num_steps
hmc_state = hmc_state._replace(num_steps=hmc_state_new.num_steps)
# reset the trajectory length
hmc_state_new = hmc_state_new._replace(
trajectory_length=hmc_state.trajectory_length)
hmc_state, z_discrete = cond(
random.bernoulli(rng_key, accept_prob),
(hmc_state_new, z_discrete_new),
identity,
(hmc_state, z_discrete),
identity,
)
# perform hmc adapting (similar to the implementation in hmc)
adapt_state = cond(
hmc_state.i < self._num_warmup,
(hmc_state.i, accept_prob, (hmc_state.z, ), hmc_state.adapt_state),
lambda args: self._wa_update(*args),
hmc_state.adapt_state,
identity,
)
itr = hmc_state.i + 1
n = jnp.where(hmc_state.i < self._num_warmup, itr,
itr - self._num_warmup)
mean_accept_prob_prev = state.hmc_state.mean_accept_prob
mean_accept_prob = (mean_accept_prob_prev +
(accept_prob - mean_accept_prob_prev) / n)
hmc_state = hmc_state._replace(
i=itr,
accept_prob=accept_prob,
mean_accept_prob=mean_accept_prob,
adapt_state=adapt_state,
)
z = {**z_discrete, **hmc_state.z}
return MixedHMCState(z, hmc_state, rng_key, accept_prob)