def test_works_for_nested_namedtuple(self):
Results = collections.namedtuple('Results', ['field1', 'inner']) # pylint: disable=invalid-name
InnerResults = collections.namedtuple('InnerResults',
['fieldA', 'fieldB']) # pylint: disable=invalid-name
accepted = Results(field1=np.int32([1, 3]),
inner=InnerResults(fieldA=np.float32([5, 7]),
fieldB=[
np.float32([9, 11]),
np.float64([13, 15]),
]))
rejected = Results(field1=np.int32([0, 2]),
inner=InnerResults(fieldA=np.float32([4, 6]),
fieldB=[
np.float32([8, 10]),
np.float64([12, 14]),
]))
chosen = choose(tf.constant([False, True]), accepted, rejected)
chosen_ = self.evaluate(chosen)
# Lhs should be 0,4,8,12 and rhs=lhs+3.
expected = Results(field1=np.int32([0, 3]),
inner=InnerResults(fieldA=np.float32([4, 7]),
fieldB=[
np.float32([8, 11]),
np.float64([12, 15]),
]))
self.assertAllClose(expected, chosen_, atol=0., rtol=1e-5)
def test_works_for_nested_namedtuple(self):
Results = collections.namedtuple('Results', ['field1', 'inner']) # pylint: disable=invalid-name
InnerResults = collections.namedtuple('InnerResults', ['fieldA', 'fieldB']) # pylint: disable=invalid-name
accepted = Results(
field1=np.int32([1, 3]),
inner=InnerResults(
fieldA=np.float32([5, 7]),
fieldB=[
np.float32([9, 11]),
np.float64([13, 15]),
]))
rejected = Results(
field1=np.int32([0, 2]),
inner=InnerResults(
fieldA=np.float32([4, 6]),
fieldB=[
np.float32([8, 10]),
np.float64([12, 14]),
]))
chosen = choose(
tf.constant([False, True]),
accepted,
rejected)
chosen_ = self.evaluate(chosen)
# Lhs should be 0,4,8,12 and rhs=lhs+3.
expected = Results(
field1=np.int32([0, 3]),
inner=InnerResults(
fieldA=np.float32([4, 7]),
fieldB=[
np.float32([8, 11]),
np.float64([12, 15]),
]))
self.assertAllClose(expected, chosen_, atol=0., rtol=1e-5)
def test_selects_batch_members_from_list_of_arrays(self):
# Shape of each array: [2, 3] = [batch_size, event_size]
# This test verifies that is_accepted selects batch members, despite the
# "usual" broadcasting being applied on the right first (event first).
zeros_states = [np.zeros((2, 3))]
ones_states = [np.ones((2, 3))]
chosen = choose(tf.constant([True, False]), zeros_states, ones_states)
chosen_ = self.evaluate(chosen)
# Make sure outer list wasn't interpreted as a dimenion of an array.
self.assertIsInstance(chosen_, list)
expected_array = np.array([
[0., 0., 0.], # zeros_states selected for first batch
[1., 1., 1.], # ones_states selected for second
])
expected = [expected_array]
self.assertAllEqual(expected, chosen_)
def one_step(self, current_state, previous_kernel_results):
"""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`).
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".
"""
with tf.name_scope(
name=mcmc_util.make_name(self.name, 'mh', 'one_step'),
values=[current_state, previous_kernel_results]):
# Take one inner step.
[
proposed_state,
proposed_results,
] = self.inner_kernel.one_step(
current_state,
previous_kernel_results.accepted_results)
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.log(tf.random_uniform(
shape=tf.shape(proposed_results.target_log_prob),
dtype=proposed_results.target_log_prob.dtype.base_dtype,
seed=self._seed_stream()))
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,
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=[],
)
return next_state, kernel_results
def one_step(self, current_state, previous_kernel_results):
"""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`).
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".
"""
# Take one inner step.
[
proposed_state,
proposed_results,
] = self.inner_kernel.one_step(
current_state, previous_kernel_results.accepted_results)
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:
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
# Note:
# - We mutate seed state so subsequent calls are not correlated.
# - We mutate seed BEFORE using it just in case users supplied the
# same seed to the inner kernel.
self._seed = distributions_util.gen_new_seed(
self.seed, salt='metropolis_hastings_one_step')
log_uniform = tf.log(
tf.random_uniform(
shape=tf.shape(proposed_results.target_log_prob),
dtype=proposed_results.target_log_prob.dtype.base_dtype,
seed=self.seed))
is_accepted = log_uniform < log_accept_ratio
independent_chain_ndims = distributions_util.prefer_static_rank(
proposed_results.target_log_prob)
next_state = mcmc_util.choose(is_accepted, proposed_state,
current_state, independent_chain_ndims)
accepted_results = type(proposed_results)(
**dict([(fn,
mcmc_util.choose(
is_accepted, getattr(proposed_results, fn),
getattr(previous_kernel_results.accepted_results, fn),
independent_chain_ndims))
for fn in proposed_results._fields]))
return [
next_state,
MetropolisHastingsKernelResults(
accepted_results=accepted_results,
is_accepted=is_accepted,
log_accept_ratio=log_accept_ratio,
proposed_state=proposed_state,
proposed_results=proposed_results,
)
]
def _swap(is_exchange_accepted, x, y):
"""Swap batches of x, y where accepted."""
with tf.name_scope('swap_where_exchange_accepted'):
new_x = mcmc_util.choose(is_exchange_accepted, y, x)
new_y = mcmc_util.choose(is_exchange_accepted, x, y)
return new_x, new_y
def one_step(self, current_state, previous_kernel_results):
"""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`).
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".
"""
with tf.name_scope(
name=mcmc_util.make_name(self.name, 'mh', 'one_step'),
values=[current_state, previous_kernel_results]):
# Take one inner step.
[
proposed_state,
proposed_results,
] = self.inner_kernel.one_step(
current_state,
previous_kernel_results.accepted_results)
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.log(tf.random_uniform(
shape=tf.shape(proposed_results.target_log_prob),
dtype=proposed_results.target_log_prob.dtype.base_dtype,
seed=self._seed_stream()))
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,
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=[],
)
return next_state, kernel_results
