def adjacent_swaps(num_replica,
batch_shape=(),
step_count=None,
seed=None):
"""Make random shuffle using only one time swaps."""
del step_count # Unused for this function.
with tf.name_scope(name or 'adjacent_swaps'):
parity_seed, proposal_seed = samplers.split_seed(seed)
# u selects parity. E.g.,
# u==False ==> [1, 0, 3, 2, 4] even parity swaps
# u==True ==> [0, 2, 1, 4, 3] odd parity swaps
# If there are only 2 replicas, then the "True" swaps are null
# swaps...which would contradict the user provided `prob_swap`.
# So special case num_replica==2, forcing u==False in this case.
u_shape = ps.concat(
(ps.ones(1, dtype=tf.int32), ps.cast(batch_shape, tf.int32)),
axis=0)
u = samplers.uniform(u_shape, seed=parity_seed) < 0.5
u = tf.where(num_replica > 2, u, False)
x = bu.left_justified_expand_dims_to(ps.range(num_replica,
dtype=tf.int64),
rank=ps.size(u_shape))
y = tf.where(tf.equal(x % 2, tf.cast(u, dtype=tf.int64)), x + 1,
x - 1)
y = tf.clip_by_value(y, 0, num_replica - 1)
# TODO(b/142689785): Consider using tf.cond and returning an empty list
# then in REMC consider using a tf.cond for short-circuiting.
return tf.where(
samplers.uniform(batch_shape, seed=proposal_seed) < prob_swap,
y, x)
def even_odd_swaps(num_replica,
batch_shape=(),
step_count=None,
seed=None):
"""Make deterministic even_odd one time swaps."""
if step_count is None:
raise ValueError('`step_count` must be supplied. Found `None`.')
del seed # Unused for this function.
with tf.name_scope(name or 'even_odd_swaps'):
# Period is 1 / frequency, and we want period = Inf if frequency = 0.
# safe_swap_period is the correct swap period in case swap_frequency > 0.
# If swap_frequency == 0, safe_swap_period is set to 1 (to avoid integer
# div by zero below). We will hard-set this case to "null swap."
swap_freq = tf.convert_to_tensor(swap_frequency,
name='swap_frequency')
safe_swap_period = tf.cast(
tf.where(swap_freq > 0,
tf.math.ceil(tf.math.reciprocal_no_nan(swap_freq)),
1),
# Although period = 1 / frequency may have roundoff error, and result
# in a period different than what the user intended, the
# user will end up with a single integer period, and thus well defined
# deterministic swaps.
tf.int32,
)
# u selects parity. E.g.,
# u==False ==> [1, 0, 3, 2, 4] even parity swaps
# u==True ==> [0, 2, 1, 4, 3] odd parity swaps
# If there are 2 replicas, then the "True" swaps are null
# swaps...which would contradict the user provided `swap_frequency`.
# So special case num_replica==2, forcing u==False in this case.
u_shape = ps.concat(
(ps.ones(1, dtype=tf.int32), ps.cast(batch_shape, tf.int32)),
axis=0)
u = tf.fill(u_shape,
tf.cast((step_count // safe_swap_period) % 2, tf.bool))
u = tf.where(num_replica > 2, u, False)
x = bu.left_justified_expand_dims_to(tf.range(num_replica,
dtype=tf.int64),
rank=ps.size(u_shape))
y = tf.where(tf.equal(x % 2, tf.cast(u, dtype=tf.int64)), x + 1,
x - 1)
y = tf.clip_by_value(y, 0, num_replica - 1)
# TODO(b/142689785): Consider using tf.cond and returning an empty list
# then in REMC consider using a tf.cond for short-circuiting.
return tf.where(
(tf.cast(step_count % safe_swap_period, tf.bool)
| tf.math.equal(swap_freq, 0)),
x, # Don't swap
y, # Swap
)
def one_step(self, current_state, previous_kernel_results, seed=None):
"""Takes one step of the TransitionKernel.
Args:
current_state: `Tensor` or Python `list` of `Tensor`s representing the
current state(s) of the Markov chain(s).
previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
`list` of `Tensor`s representing internal calculations made within the
previous call to this function (or as returned by `bootstrap_results`).
seed: Optional, a seed for reproducible sampling.
Returns:
next_state: `Tensor` or Python `list` of `Tensor`s representing the
next state(s) of the Markov chain(s).
kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
`Tensor`s representing internal calculations made within this function.
This inculdes replica states.
"""
# The code below propagates one step states of shape
# [n_replica] + batch_shape + event_shape.
#
# The step is done in three parts:
# 1) Call one_step to transition states via a tempered version of
# self.target_log_prob_fn (see _replica_target_log_prob).
# 2) Permute values in states
# 3) Update state-dependent values, such as log_probs.
#
# We chose to swap states, rather than temperatures, because...
# (i) If swapping temperatures, you *still* have to swap log_probs to
# determine acceptance, as well as states (for kernel results).
# So it's just as difficult to swap temperatures.
# (ii) If swapping temperatures, you have to take care to swap any user-
# supplied temperature related things (like step size).
# A-priori, we don't know what else will need to be swapped!
# (iii)In both cases, the kernel results need to be updated in a non-trivial
# manner....so we either special-case, or use bootstrap.
with tf.name_scope(mcmc_util.make_name(self.name, 'remc', 'one_step')):
# Force a read in case the `inverse_temperatures` is a `tf.Variable`.
inverse_temperatures = tf.convert_to_tensor(
previous_kernel_results.inverse_temperatures,
name='inverse_temperatures')
target_log_prob_for_inner_kernel = _make_replica_target_log_prob_fn(
target_log_prob_fn=self.target_log_prob_fn,
inverse_temperatures=inverse_temperatures,
untempered_log_prob_fn=self.untempered_log_prob_fn,
tempered_log_prob_fn=self.tempered_log_prob_fn,
)
# TODO(b/159636942): Clean up the helpful error msg after 2020-11-10.
try:
inner_kernel = self.make_kernel_fn( # pylint: disable=not-callable
target_log_prob_for_inner_kernel)
except TypeError as e:
if 'argument' not in str(e):
raise
raise TypeError(
'`ReplicaExchangeMC`s `make_kernel_fn` no longer receives a `seed` '
'argument. `TransitionKernel` instances now receive seeds via '
'`one_step`.')
seed = samplers.sanitize_seed(seed) # Retain for diagnostics.
inner_seed, swap_seed, logu_seed = samplers.split_seed(seed, n=3)
# Step the inner TransitionKernel.
[
pre_swap_replica_states,
pre_swap_replica_results,
] = inner_kernel.one_step(
previous_kernel_results.post_swap_replica_states,
previous_kernel_results.post_swap_replica_results,
seed=inner_seed)
pre_swap_replica_target_log_prob = _get_field(
# These are tempered log probs (have been divided by temperature).
pre_swap_replica_results,
'target_log_prob')
dtype = pre_swap_replica_target_log_prob.dtype
replica_and_batch_shape = ps.shape(
pre_swap_replica_target_log_prob)
batch_shape = replica_and_batch_shape[1:]
replica_and_batch_rank = ps.rank(pre_swap_replica_target_log_prob)
num_replica = ps.size0(inverse_temperatures)
inverse_temperatures = bu.left_justified_broadcast_to(
inverse_temperatures, replica_and_batch_shape)
# Now that each replica has done one_step, it is time to consider swaps.
# swap.shape = [n_replica], and is a "once only" permutation, meaning it
# is achievable by a sequence of pairwise permutations, where each element
# is moved at most once.
# E.g. if swaps = [1, 0, 2], we will consider swapping temperatures 0 and
# 1, keeping 2 fixed. This exact same swap is considered for *every*
# batch member. Of course some batch members may accept and some reject.
try:
swaps = tf.cast(
self.swap_proposal_fn( # pylint: disable=not-callable
num_replica,
batch_shape=batch_shape,
seed=swap_seed,
step_count=previous_kernel_results.step_count),
dtype=tf.int32)
except TypeError as e:
if 'step_count' not in str(e):
raise
warnings.warn(
'The `swap_proposal_fn` given to ReplicaExchangeMC did not accept '
'the `step_count` argument. Falling back to omitting the '
'argument. This fallback will be removed after 24-Oct-2020.'
)
swaps = tf.cast(
self.swap_proposal_fn( # pylint: disable=not-callable
num_replica,
batch_shape=batch_shape,
seed=swap_seed),
dtype=tf.int32)
null_swaps = bu.left_justified_expand_dims_like(
tf.range(num_replica, dtype=swaps.dtype), swaps)
swaps = _maybe_embed_swaps_validation(swaps, null_swaps,
self.validate_args)
# Un-temper the log probs for use in the swap acceptance ratio.
if self.tempered_log_prob_fn is None:
# Efficient way of re-evaluating target_log_prob_fn on the
# pre_swap_replica_states.
untempered_negative_energy_ignoring_ulp = (
# Since untempered_log_prob_fn is None, we may assume
# inverse_temperatures > 0 (else the target is improper).
pre_swap_replica_target_log_prob / inverse_temperatures)
else:
# The untempered_log_prob_fn does not factor into the acceptance ratio.
# Proof: Suppose the tempered target is
# p_k(x) = f(x)^{beta_k} g(x),
# So f(x) is tempered, and g(x) is not. Then, the acceptance ratio for
# a 1 <--> 2 swap is...
# (p_1(x_2) p_2(x_1)) / (p_1(x_1) p_2(x_2))
# which depends only on f(x), since terms involving g(x) cancel.
untempered_negative_energy_ignoring_ulp = self.tempered_log_prob_fn(
*pre_swap_replica_states)
# Since `swaps` is its own inverse permutation we automatically know the
# swap counterpart: range(num_replica). We use this idea to compute the
# acceptance in a vectorized manner at the cost of wasting roughly half
# our computation. Although we could use `unique` to solve this problem,
# we expect the cost of `unique` to be higher than the dozens of wasted
# arithmetic calculations. Worse, it'd mean we need dynamic sized Tensors
# (eg, using `tf.where(bool)`) and so we wouldn't be able to XLA compile.
# Note: diffs would normally be "proposed - current" however energy is
# flipped since `energy == -log_prob`.
# Note: The untempered_log_prob_fn (if provided) is not included in
# untempered_pre_swap_replica_target_log_prob, and hence does not factor
# into energy_diff. Why? Because, it cancels out in the acceptance ratio.
energy_diff = (untempered_negative_energy_ignoring_ulp -
mcmc_util.index_remapping_gather(
untempered_negative_energy_ignoring_ulp,
swaps,
name='gather_swap_tlp'))
swapped_inverse_temperatures = mcmc_util.index_remapping_gather(
inverse_temperatures, swaps, name='gather_swap_temps')
inverse_temp_diff = swapped_inverse_temperatures - inverse_temperatures
# If i and j are swapping, log_accept_ratio[] i and j are equal.
log_accept_ratio = (energy_diff * bu.left_justified_expand_dims_to(
inverse_temp_diff, replica_and_batch_rank))
log_accept_ratio = tf.where(tf.math.is_finite(log_accept_ratio),
log_accept_ratio,
tf.constant(-np.inf, dtype=dtype))
# Produce log[Uniform] draws that are identical at swapped indices.
log_uniform = tf.math.log(
samplers.uniform(shape=replica_and_batch_shape,
dtype=dtype,
seed=logu_seed))
anchor_swaps = tf.minimum(swaps, null_swaps)
log_uniform = mcmc_util.index_remapping_gather(
log_uniform, anchor_swaps)
is_swap_accepted_mask = tf.less(log_uniform,
log_accept_ratio,
name='is_swap_accepted_mask')
def _swap_tensor(x):
return mcmc_util.choose(
is_swap_accepted_mask,
mcmc_util.index_remapping_gather(x, swaps), x)
post_swap_replica_states = [
_swap_tensor(s) for s in pre_swap_replica_states
]
expanded_null_swaps = bu.left_justified_broadcast_to(
null_swaps, replica_and_batch_shape)
is_swap_proposed = _compute_swap_notmatrix(
# Broadcast both so they have shape [num_replica] + batch_shape.
# This (i) makes them have same shape as is_swap_accepted, and
# (ii) keeps shape consistent if someday swaps has a batch shape.
expanded_null_swaps,
bu.left_justified_broadcast_to(swaps, replica_and_batch_shape))
# To get is_swap_accepted in ordered position, we use
# _compute_swap_notmatrix on current and next replica positions.
post_swap_replica_position = _swap_tensor(expanded_null_swaps)
is_swap_accepted = _compute_swap_notmatrix(
post_swap_replica_position, expanded_null_swaps)
if self._state_includes_replicas:
post_swap_states = post_swap_replica_states
else:
post_swap_states = [s[0] for s in post_swap_replica_states]
post_swap_replica_results = _set_swapped_fields_to_nan(
_swap_log_prob_and_maybe_grads(pre_swap_replica_results,
post_swap_replica_states,
inner_kernel))
if mcmc_util.is_list_like(current_state):
# We *always* canonicalize the states in the kernel results.
states = post_swap_states
else:
states = post_swap_states[0]
post_swap_kernel_results = ReplicaExchangeMCKernelResults(
post_swap_replica_states=post_swap_replica_states,
pre_swap_replica_results=pre_swap_replica_results,
post_swap_replica_results=post_swap_replica_results,
is_swap_proposed=is_swap_proposed,
is_swap_accepted=is_swap_accepted,
is_swap_proposed_adjacent=_sub_diag(is_swap_proposed),
is_swap_accepted_adjacent=_sub_diag(is_swap_accepted),
# Store the original pkr.inverse_temperatures in case its a
# `tf.Variable`.
inverse_temperatures=previous_kernel_results.
inverse_temperatures,
swaps=swaps,
step_count=previous_kernel_results.step_count + 1,
seed=seed,
potential_energy=-untempered_negative_energy_ignoring_ulp,
)
return states, post_swap_kernel_results
def checkAllIntegrals(self, prior_scale, likelihood_scale,
inverse_temperatures, iid_chain_ndims):
prior_scale = tf.convert_to_tensor(prior_scale, name='prior_scale')
likelihood_scale = tf.convert_to_tensor(likelihood_scale,
name='likelihood_scale')
# Create (normalized) prior and likelihood. Their product of course is not
# normalized. In particular, there is a number `normalizing_const` such that
# posterior(z) = prior.prob(x) * likelihood.prob(x) / normalizing_const
# is normalized.
prior = tfd.Normal(0., prior_scale)
likelihood = tfd.Normal(0., likelihood_scale)
posterior = tfd.Normal(0., (prior_scale**-2 +
likelihood_scale**-2)**(-0.5))
# Get a good step size, custom for every replica/batch member.
bcast_inv_temperatures = bu.left_justified_expand_dims_to(
inverse_temperatures,
# Broadcast over replicas.
1 +
# Broadcast over chains.
iid_chain_ndims +
# Broadcast over batch dims.
tf.rank(likelihood_scale))
tempered_posteriors = tfd.Normal(
0.,
# One tempered posterior for every inverse_temperature.
(prior_scale**-2 + bcast_inv_temperatures * likelihood_scale**-2
)**(-0.5))
step_size = 0.71234 * tempered_posteriors.stddev()
num_leapfrog_steps = tf.cast(
tf.math.ceil(1.567 / tf.reduce_min(step_size)), tf.int32)
def make_kernel_fn(target_log_prob_fn):
return tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=target_log_prob_fn,
step_size=step_size,
num_leapfrog_steps=num_leapfrog_steps,
)
remc = tfp.mcmc.ReplicaExchangeMC(
target_log_prob_fn=None,
untempered_log_prob_fn=prior.log_prob,
tempered_log_prob_fn=likelihood.log_prob,
inverse_temperatures=inverse_temperatures,
state_includes_replicas=False,
make_kernel_fn=make_kernel_fn,
swap_proposal_fn=tfp.mcmc.even_odd_swap_proposal_fn(1.),
)
def trace_fn(state, results): # pylint: disable=unused-argument
return {
'replica_log_accept_ratio':
results.post_swap_replica_results.log_accept_ratio,
'is_swap_accepted_adjacent': results.is_swap_accepted_adjacent,
'is_swap_proposed_adjacent': results.is_swap_proposed_adjacent,
'potential_energy': results.potential_energy,
}
if tf.executing_eagerly():
num_results = 100
else:
num_results = 1000
num_burnin_steps = num_results // 10
n_samples_per_chain = 2
initial_sample_shape = [n_samples_per_chain] * iid_chain_ndims
unused_replica_states, trace = self.evaluate(
tfp.mcmc.sample_chain(
num_results=num_results,
# Start at one of the modes, in order to make mode jumping necessary
# if we want to pass test.
current_state=prior.sample(initial_sample_shape,
seed=test_util.test_seed()),
kernel=remc,
num_burnin_steps=num_burnin_steps,
trace_fn=trace_fn,
seed=test_util.test_seed()))
# Tolerance depends on samples * replicas * number of (iid) chains.
# ess.shape = [n_replica, ...]
# We will sum over batch dims, then take min over replica.
ess = tfp.mcmc.effective_sample_size(trace['potential_energy'])
if iid_chain_ndims:
ess = tf.reduce_sum(ess, axis=tf.range(1, 1 + iid_chain_ndims))
min_ess = self.evaluate(tf.reduce_min(ess))
n_combined_results = min_ess * inverse_temperatures.shape[0]
# Make sure sampling worked well enough, for every replica/chain.
conditional_swap_prob = (
np.sum(trace['is_swap_accepted_adjacent'], axis=0) /
np.sum(trace['is_swap_proposed_adjacent'], axis=0))
self.assertAllGreater(conditional_swap_prob, 0.5)
replica_mean_accept_prob = np.mean(np.exp(
np.minimum(0, trace['replica_log_accept_ratio'])),
axis=0)
self.assertAllGreater(replica_mean_accept_prob, 0.5)
integrals = self.evaluate(
tfp.experimental.mcmc.remc_thermodynamic_integrals(
inverse_temperatures,
trace['potential_energy'],
iid_chain_ndims=iid_chain_ndims,
))
self.assertAllEqual(posterior.batch_shape,
integrals.log_normalizing_constant_ratio.shape)
actual_log_normalizing_const = self.evaluate(
# Use arbitrary point, 0, to find the constant.
prior.log_prob(0.) + likelihood.log_prob(0.) -
posterior.log_prob(0.))
self.assertAllClose(integrals.log_normalizing_constant_ratio,
actual_log_normalizing_const,
rtol=10 / np.sqrt(n_combined_results))
self.assertAllEqual(posterior.batch_shape,
integrals.cross_entropy_difference.shape)
def cross_entropy(dist):
z = dist.sample(50000, seed=test_util.test_seed())
return tf.reduce_mean(likelihood.log_prob(z), axis=0)
iid_cross_entropy_difference = self.evaluate(
cross_entropy(posterior) - cross_entropy(prior))
self.assertAllClose(integrals.cross_entropy_difference,
iid_cross_entropy_difference,
rtol=30 / np.sqrt(n_combined_results))
