def bootstrap_results(self, init_state):
"""Returns an object with the same type as returned by `one_step`.
Args:
init_state: `Tensor` or Python `list` of `Tensor`s representing the
initial state(s) of the Markov chain(s).
Returns:
kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
`Tensor`s representing internal calculations made within this function.
Raises:
ValueError: if `inner_kernel` results doesn't contain the member
"target_log_prob".
"""
with tf.name_scope(
mcmc_util.make_name(self.name, 'mh', 'bootstrap_results')):
pkr = self.inner_kernel.bootstrap_results(init_state)
if not has_target_log_prob(pkr):
raise ValueError(
'"target_log_prob" must be a member of `inner_kernel` results.'
)
x = pkr.target_log_prob
return MetropolisHastingsKernelResults(
# See note regarding `strip_seeds` above in `one_step`.
accepted_results=mcmc_util.strip_seeds(pkr),
is_accepted=tf.ones_like(x, dtype=tf.bool),
log_accept_ratio=tf.zeros_like(x),
proposed_state=init_state,
proposed_results=pkr,
extra=[],
# Allow room for one_step's seed.
seed=samplers.zeros_seed(),
)
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: PRNG seed; see `tfp.random.sanitize_seed` for details.
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.
Raises:
ValueError: if `inner_kernel` results doesn't contain the member
"target_log_prob".
"""
is_seeded = seed is not None
seed = samplers.sanitize_seed(seed) # Retain for diagnostics.
proposal_seed, acceptance_seed = samplers.split_seed(seed)
with tf.name_scope(mcmc_util.make_name(self.name, 'mh', 'one_step')):
# Take one inner step.
inner_kwargs = dict(seed=proposal_seed) if is_seeded else {}
[
proposed_state,
proposed_results,
] = self.inner_kernel.one_step(
current_state, previous_kernel_results.accepted_results,
**inner_kwargs)
if mcmc_util.is_list_like(current_state):
proposed_state = tf.nest.pack_sequence_as(
current_state, proposed_state)
if (not has_target_log_prob(proposed_results)
or not has_target_log_prob(
previous_kernel_results.accepted_results)):
raise ValueError('"target_log_prob" must be a member of '
'`inner_kernel` results.')
# Compute log(acceptance_ratio).
to_sum = [
proposed_results.target_log_prob,
-previous_kernel_results.accepted_results.target_log_prob
]
try:
if (not mcmc_util.is_list_like(
proposed_results.log_acceptance_correction)
or proposed_results.log_acceptance_correction):
to_sum.append(proposed_results.log_acceptance_correction)
except AttributeError:
warnings.warn(
'Supplied inner `TransitionKernel` does not have a '
'`log_acceptance_correction`. Assuming its value is `0.`')
log_accept_ratio = mcmc_util.safe_sum(
to_sum, name='compute_log_accept_ratio')
# If proposed state reduces likelihood: randomly accept.
# If proposed state increases likelihood: always accept.
# I.e., u < min(1, accept_ratio), where u ~ Uniform[0,1)
# ==> log(u) < log_accept_ratio
log_uniform = tf.math.log(
samplers.uniform(shape=prefer_static.shape(
proposed_results.target_log_prob),
dtype=dtype_util.base_dtype(
proposed_results.target_log_prob.dtype),
seed=acceptance_seed))
is_accepted = log_uniform < log_accept_ratio
next_state = mcmc_util.choose(is_accepted,
proposed_state,
current_state,
name='choose_next_state')
kernel_results = MetropolisHastingsKernelResults(
accepted_results=mcmc_util.choose(
is_accepted,
# We strip seeds when populating `accepted_results` because unlike
# other kernel result fields, seeds are not a per-chain value.
# Thus it is impossible to choose between a previously accepted
# seed value and a proposed seed, since said choice would need to
# be made on a per-chain basis.
mcmc_util.strip_seeds(proposed_results),
previous_kernel_results.accepted_results,
name='choose_inner_results'),
is_accepted=is_accepted,
log_accept_ratio=log_accept_ratio,
proposed_state=proposed_state,
proposed_results=proposed_results,
extra=[],
seed=seed,
)
return next_state, kernel_results