def testNoParameters(self):
kernel = FakeInnerNoParameters()
new_kernel = util.enable_store_parameters_in_results(kernel)
self.assertIs(kernel, new_kernel)
def __init__(self,
inner_kernel,
num_adaptation_steps,
target_accept_prob=0.75,
adaptation_rate=0.01,
step_size_setter_fn=_hmc_like_step_size_setter_fn,
step_size_getter_fn=_hmc_like_step_size_getter_fn,
log_accept_prob_getter_fn=_hmc_like_log_accept_prob_getter_fn,
validate_args=False,
name=None):
"""Creates the step size adaptation kernel.
The default setter_fn and the getter_fn callbacks assume that the inner
kernel produces kernel results structurally the same as the
`HamiltonianMonteCarlo` kernel.
Args:
inner_kernel: `TransitionKernel`-like object.
num_adaptation_steps: Scalar `int` `Tensor` number of initial steps to
during which to adjust the step size. This may be greater, less than, or
equal to the number of burnin steps.
target_accept_prob: A floating point `Tensor` representing desired
acceptance probability. Must be a positive number less than 1. This can
either be a scalar, or have shape [num_chains]. Default value: `0.75`
(the [center of asymptotically optimal rate for HMC][1]).
adaptation_rate: `Tensor` representing amount to scale the current
`step_size`.
step_size_setter_fn: A callable with the signature
`(kernel_results, new_step_size) -> new_kernel_results` where
`kernel_results` are the results of the `inner_kernel`, `new_step_size`
is a `Tensor` or a nested collection of `Tensor`s with the same
structure as returned by the `step_size_getter_fn`, and
`new_kernel_results` are a copy of `kernel_results` with the step
size(s) set.
step_size_getter_fn: A callable with the signature
`(kernel_results) -> step_size` where `kernel_results` are the results
of the `inner_kernel`, and `step_size` is a floating point `Tensor` or a
nested collection of such `Tensor`s.
log_accept_prob_getter_fn: A callable with the signature
`(kernel_results) -> log_accept_prob` where `kernel_results` are the
results of the `inner_kernel`, and `log_accept_prob` is a floating point
`Tensor`. `log_accept_prob` can either be a scalar, or have shape
[num_chains]. If it's the latter, `step_size` should also have the same
leading dimension.
validate_args: Python `bool`. When `True` kernel parameters are checked
for validity. When `False` invalid inputs may silently render incorrect
outputs.
name: Python `str` name prefixed to Ops created by this class. Default:
'simple_step_size_adaptation'.
#### References
[1]: Betancourt, M. J., Byrne, S., & Girolami, M. (2014). _Optimizing The
Integrator Step Size for Hamiltonian Monte Carlo_.
http://arxiv.org/abs/1411.6669
"""
inner_kernel = mcmc_util.enable_store_parameters_in_results(
inner_kernel)
with tf.compat.v1.name_scope(mcmc_util.make_name(
name, 'simple_step_size_adaptation', '__init__'),
values=[
target_accept_prob, adaptation_rate,
num_adaptation_steps
]) as name:
dtype = dtype_util.common_dtype(
[target_accept_prob, adaptation_rate], tf.float32)
target_accept_prob = tf.convert_to_tensor(
value=target_accept_prob,
dtype=dtype,
name='target_accept_prob')
adaptation_rate = tf.convert_to_tensor(value=adaptation_rate,
dtype=dtype,
name='adaptation_rate')
num_adaptation_steps = tf.convert_to_tensor(
value=num_adaptation_steps,
dtype=tf.int32,
name='num_adaptation_steps')
target_accept_prob = _maybe_validate_target_accept_prob(
target_accept_prob, validate_args)
self._parameters = dict(
inner_kernel=inner_kernel,
num_adaptation_steps=num_adaptation_steps,
target_accept_prob=target_accept_prob,
adaptation_rate=adaptation_rate,
step_size_setter_fn=step_size_setter_fn,
step_size_getter_fn=step_size_getter_fn,
log_accept_prob_getter_fn=log_accept_prob_getter_fn,
name=name,
)
def __init__(
self,
inner_kernel,
num_adaptation_steps,
use_halton_sequence_jitter=True,
adaptation_rate=0.025,
jitter_amount=1.,
criterion_fn=chees_criterion,
max_leapfrog_steps=1000,
num_leapfrog_steps_getter_fn=hmc_like_num_leapfrog_steps_getter_fn,
num_leapfrog_steps_setter_fn=hmc_like_num_leapfrog_steps_setter_fn,
step_size_getter_fn=hmc_like_step_size_getter_fn,
proposed_velocity_getter_fn=hmc_like_proposed_velocity_getter_fn,
log_accept_prob_getter_fn=hmc_like_log_accept_prob_getter_fn,
proposed_state_getter_fn=hmc_like_proposed_state_getter_fn,
validate_args=False,
name=None):
"""Creates the trajectory length adaptation kernel.
The default setter_fn and the getter_fn callbacks assume that the inner
kernel produces kernel results structurally the same as the
`HamiltonianMonteCarlo` kernel (possibly wrapped in some step size
adaptation kernel).
Args:
inner_kernel: `TransitionKernel`-like object.
num_adaptation_steps: Scalar `int` `Tensor` number of initial steps to
during which to adjust the step size. This may be greater, less than, or
equal to the number of burnin steps.
use_halton_sequence_jitter: Python bool. Whether to use a Halton sequence
for jittering the trajectory length. This makes the procedure more
stable than sampling trajectory lengths from a uniform distribution.
adaptation_rate: Floating point scalar `Tensor`. How rapidly to adapt the
trajectory length.
jitter_amount: Floating point scalar `Tensor`. How much to jitter the
trajectory on the next step. The trajectory length is sampled from `[(1
- jitter_amount) * max_trajectory_length, max_trajectory_length]`.
criterion_fn: Callable with `(previous_state, proposed_state, accept_prob)
-> criterion`. Computes the criterion value.
max_leapfrog_steps: Int32 scalar `Tensor`. Clips the number of leapfrog
steps to this value.
num_leapfrog_steps_getter_fn: A callable with the signature
`(kernel_results) -> num_leapfrog_steps` where `kernel_results` are the
results of the `inner_kernel`, and `num_leapfrog_steps` is a floating
point `Tensor`.
num_leapfrog_steps_setter_fn: A callable with the signature
`(kernel_results, new_num_leapfrog_steps) -> new_kernel_results` where
`kernel_results` are the results of the `inner_kernel`,
`new_num_leapfrog_steps` is a scalar tensor `Tensor`, and
`new_kernel_results` are a copy of `kernel_results` with the number of
leapfrog steps set.
step_size_getter_fn: A callable with the signature `(kernel_results) ->
step_size` where `kernel_results` are the results of the `inner_kernel`,
and `step_size` is a floating point `Tensor`.
proposed_velocity_getter_fn: A callable with the signature
`(kernel_results) -> proposed_velocity` where `kernel_results` are the
results of the `inner_kernel`, and `proposed_velocity` is a (possibly
nested) floating point `Tensor`. Velocity is derivative of state with
respect to trajectory length.
log_accept_prob_getter_fn: A callable with the signature `(kernel_results)
-> log_accept_prob` where `kernel_results` are the results of the
`inner_kernel`, and `log_accept_prob` is a floating point `Tensor`.
`log_accept_prob` has shape `[C0, ...., Cb]` with `b > 0`.
proposed_state_getter_fn: A callable with the signature `(kernel_results)
-> proposed_state` where `kernel_results` are the results of the
`inner_kernel`, and `proposed_state` is a (possibly nested) floating
point `Tensor`.
validate_args: Python `bool`. When `True` kernel parameters are checked
for validity. When `False` invalid inputs may silently render incorrect
outputs.
name: Python `str` name prefixed to Ops created by this class. Default:
'simple_step_size_adaptation'.
Raises:
ValueError: If `inner_kernel` contains a `TransformedTransitionKernel` in
its hierarchy. If you need to use the `TransformedTransitionKernel`,
place it above this kernel in the hierarchy (see the example in the
class docstring).
"""
inner_kernel = mcmc_util.enable_store_parameters_in_results(inner_kernel)
_forbid_inner_transformed_kernel(inner_kernel)
with tf.name_scope(
mcmc_util.make_name(name, 'gradient_based_trajectory_length_adaptation',
'__init__')) as name:
dtype = dtype_util.common_dtype([adaptation_rate, jitter_amount],
tf.float32)
num_adaptation_steps = tf.convert_to_tensor(
num_adaptation_steps, dtype=tf.int32, name='num_adaptation_steps')
adaptation_rate = tf.convert_to_tensor(
adaptation_rate, dtype=dtype, name='adaptation_rate')
jitter_amount = tf.convert_to_tensor(
jitter_amount, dtype=dtype, name='jitter_amount')
max_leapfrog_steps = tf.convert_to_tensor(
max_leapfrog_steps, dtype=tf.int32, name='max_leapfrog_steps')
self._parameters = dict(
inner_kernel=inner_kernel,
num_adaptation_steps=num_adaptation_steps,
use_halton_sequence_jitter=use_halton_sequence_jitter,
adaptation_rate=adaptation_rate,
jitter_amount=jitter_amount,
criterion_fn=criterion_fn,
max_leapfrog_steps=max_leapfrog_steps,
num_leapfrog_steps_getter_fn=num_leapfrog_steps_getter_fn,
num_leapfrog_steps_setter_fn=num_leapfrog_steps_setter_fn,
step_size_getter_fn=step_size_getter_fn,
proposed_velocity_getter_fn=proposed_velocity_getter_fn,
log_accept_prob_getter_fn=log_accept_prob_getter_fn,
proposed_state_getter_fn=hmc_like_proposed_state_getter_fn,
validate_args=validate_args,
name=name,
)
def __init__(
self,
inner_kernel,
num_adaptation_steps,
target_accept_prob=0.75,
exploration_shrinkage=0.05,
shrinkage_target=None,
step_count_smoothing=10,
decay_rate=0.75,
step_size_setter_fn=hmc_like_step_size_setter_fn,
step_size_getter_fn=hmc_like_step_size_getter_fn,
log_accept_prob_getter_fn=hmc_like_log_accept_prob_getter_fn,
reduce_fn=reduce_logmeanexp,
experimental_reduce_chain_axis_names=None,
validate_args=False,
name=None):
"""Initializes this transition kernel.
Args:
inner_kernel: `TransitionKernel`-like object.
num_adaptation_steps: Scalar `int` `Tensor` number of initial steps to
during which to adjust the step size. This may be greater, less than, or
equal to the number of burnin steps.
target_accept_prob: A floating point `Tensor` representing desired
acceptance probability. Must be a positive number less than 1. This can
either be a scalar, or have shape [num_chains]. Default value: `0.75`
(the [center of asymptotically optimal rate for HMC][1]).
exploration_shrinkage: Floating point scalar `Tensor`. How strongly the
exploration rate is biased towards the shrinkage target.
shrinkage_target: `Tensor` or list of tensors. Value the exploration
step size(s) is/are biased towards.
As `num_adaptation_steps --> infinity`, this bias goes to zero.
Defaults to 10 times the initial step size.
step_count_smoothing: Int32 scalar `Tensor`. Number of "pseudo-steps"
added to the number of steps taken to prevents noisy exploration during
the early samples.
decay_rate: Floating point scalar `Tensor`. How much to favor recent
iterations over earlier ones. A value of 1 gives equal weight to all
history. A value of 0 gives weight only to the most recent iteration.
step_size_setter_fn: A callable with the signature `(kernel_results,
new_step_size) -> new_kernel_results` where `kernel_results` are the
results of the `inner_kernel`, `new_step_size` is a `Tensor` or a nested
collection of `Tensor`s with the same structure as returned by the
`step_size_getter_fn`, and `new_kernel_results` are a copy of
`kernel_results` with the step size(s) set.
step_size_getter_fn: A callable with the signature `(kernel_results) ->
step_size` where `kernel_results` are the results of the `inner_kernel`,
and `step_size` is a floating point `Tensor` or a nested collection of
such `Tensor`s.
log_accept_prob_getter_fn: A callable with the signature `(kernel_results)
-> log_accept_prob` where `kernel_results` are the results of the
`inner_kernel`, and `log_accept_prob` is a floating point `Tensor`.
`log_accept_prob` can either be a scalar, or have shape [num_chains]. If
it's the latter, `step_size` should also have the same leading
dimension.
reduce_fn: A callable with signature `(input_tensor, axis, keepdims) ->
tensor` that returns a log-reduction of `log_accept_prob`, typically
some sort of mean. By default, this performs an arithmetic mean.
experimental_reduce_chain_axis_names: A `str` or list of `str`s indicating
the named axes that should additionally reduced during the log-reduction
of `log_accept_prob`.
validate_args: Python `bool`. When `True` kernel parameters are checked
for validity. When `False` invalid inputs may silently render incorrect
outputs.
name: Python `str` name prefixed to Ops created by this function.
Default value: `None` (i.e., 'dual_averaging_step_size_adaptation').
"""
inner_kernel = mcmc_util.enable_store_parameters_in_results(inner_kernel)
with tf.name_scope(
mcmc_util.make_name(
name, 'dual_averaging_step_size_adaptation', '__init__')) as name:
dtype = dtype_util.common_dtype([
target_accept_prob,
exploration_shrinkage,
shrinkage_target,
decay_rate
], dtype_hint=tf.float32)
target_accept_prob = tf.convert_to_tensor(
target_accept_prob, dtype=dtype, name='target_accept_prob')
exploration_shrinkage = tf.convert_to_tensor(
exploration_shrinkage, dtype=dtype, name='exploration_shrinkage')
step_count_smoothing = tf.cast(
# Cast to dtype, since we asked the user to provide Int32, but we
# want to convert to tensor here (and make it `dtype`).
# I.e., convert_to_tensor will fail here if the user did what we
# asked.
step_count_smoothing, dtype=dtype, name='step_count_smoothing')
decay_rate = tf.convert_to_tensor(
decay_rate, dtype=dtype, name='decay_rate')
num_adaptation_steps = tf.convert_to_tensor(
num_adaptation_steps, dtype=tf.int32, name='num_adaptation_steps')
target_accept_prob = _maybe_validate_target_accept_prob(
target_accept_prob, validate_args)
if shrinkage_target is not None:
def _convert(x):
return tf.convert_to_tensor(x, dtype=dtype, name='shrinkage_target')
shrinkage_target = tf.nest.map_structure(_convert, shrinkage_target)
self._parameters = dict(
inner_kernel=inner_kernel,
num_adaptation_steps=num_adaptation_steps,
target_accept_prob=target_accept_prob,
exploration_shrinkage=exploration_shrinkage,
shrinkage_target=shrinkage_target,
step_count_smoothing=step_count_smoothing,
decay_rate=decay_rate,
step_size_setter_fn=step_size_setter_fn,
step_size_getter_fn=step_size_getter_fn,
log_accept_prob_getter_fn=log_accept_prob_getter_fn,
reduce_fn=reduce_fn,
experimental_reduce_chain_axis_names=(
experimental_reduce_chain_axis_names),
name=name)
def __init__(self,
inner_kernel,
num_adaptation_steps,
target_accept_prob=0.75,
exploration_shrinkage=0.05,
step_count_smoothing=10,
decay_rate=0.75,
step_size_setter_fn=hmc_like_step_size_setter_fn,
step_size_getter_fn=hmc_like_step_size_getter_fn,
log_accept_prob_getter_fn=hmc_like_log_accept_prob_getter_fn,
validate_args=False,
name=None):
"""Initializes this transition kernel.
Args:
inner_kernel: `TransitionKernel`-like object.
num_adaptation_steps: Scalar `int` `Tensor` number of initial steps to
during which to adjust the step size. This may be greater, less than, or
equal to the number of burnin steps.
target_accept_prob: A floating point `Tensor` representing desired
acceptance probability. Must be a positive number less than 1. This can
either be a scalar, or have shape [num_chains]. Default value: `0.75`
(the [center of asymptotically optimal rate for HMC][1]).
exploration_shrinkage: Floating point scalar `Tensor`. How strongly the
exploration rate is biased towards the shrinkage target.
step_count_smoothing: Int32 scalar `Tensor`. Number of "pseudo-steps"
added to the number of steps taken to prevents noisy exploration during
the early samples.
decay_rate: Floating point scalar `Tensor`. How much to favor recent
iterations over earlier ones. A value of 1 gives equal weight to all
history.
step_size_setter_fn: A callable with the signature `(kernel_results,
new_step_size) -> new_kernel_results` where `kernel_results` are the
results of the `inner_kernel`, `new_step_size` is a `Tensor` or a nested
collection of `Tensor`s with the same structure as returned by the
`step_size_getter_fn`, and `new_kernel_results` are a copy of
`kernel_results` with the step size(s) set.
step_size_getter_fn: A callable with the signature `(kernel_results) ->
step_size` where `kernel_results` are the results of the `inner_kernel`,
and `step_size` is a floating point `Tensor` or a nested collection of
such `Tensor`s.
log_accept_prob_getter_fn: A callable with the signature `(kernel_results)
-> log_accept_prob` where `kernel_results` are the results of the
`inner_kernel`, and `log_accept_prob` is a floating point `Tensor`.
`log_accept_prob` can either be a scalar, or have shape [num_chains]. If
it's the latter, `step_size` should also have the same leading
dimension.
validate_args: Python `bool`. When `True` kernel parameters are checked
for validity. When `False` invalid inputs may silently render incorrect
outputs.
name: Python `str` name prefixed to Ops created by this function.
Default value: `None` (i.e., 'dual_averaging_step_size_adaptation').
"""
inner_kernel = mcmc_util.enable_store_parameters_in_results(
inner_kernel)
with tf.name_scope(
mcmc_util.make_name(name,
'dual_averaging_step_size_adaptation',
'__init__')) as name:
dtype = dtype_util.common_dtype(
[target_accept_prob, exploration_shrinkage, decay_rate],
dtype_hint=tf.float32)
target_accept_prob = tf.convert_to_tensor(
target_accept_prob, dtype=dtype, name='target_accept_prob')
exploration_shrinkage = tf.convert_to_tensor(
exploration_shrinkage,
dtype=dtype,
name='exploration_shrinkage')
step_count_smoothing = tf.convert_to_tensor(
step_count_smoothing, dtype=dtype, name='step_count_smoothing')
decay_rate = tf.convert_to_tensor(decay_rate,
dtype=dtype,
name='decay_rate')
num_adaptation_steps = tf.convert_to_tensor(
num_adaptation_steps,
dtype=tf.int32,
name='num_adaptation_steps')
target_accept_prob = _maybe_validate_target_accept_prob(
target_accept_prob, validate_args)
self._parameters = dict(
inner_kernel=inner_kernel,
num_adaptation_steps=num_adaptation_steps,
target_accept_prob=target_accept_prob,
exploration_shrinkage=exploration_shrinkage,
step_count_smoothing=step_count_smoothing,
decay_rate=decay_rate,
step_size_setter_fn=step_size_setter_fn,
step_size_getter_fn=step_size_getter_fn,
log_accept_prob_getter_fn=log_accept_prob_getter_fn,
name=name)
def __init__(
self,
inner_kernel,
initial_running_variance,
num_estimation_steps=None,
momentum_distribution_setter_fn=hmc_like_momentum_distribution_setter_fn,
momentum_distribution_getter_fn=hmc_like_momentum_distribution_getter_fn,
validate_args=False,
experimental_shard_axis_names=None,
name=None):
"""Creates the diagonal mass matrix adaptation kernel.
Users must provide an `initial_running_variance`, either from a previous
`DiagonalMassMatrixAdaptation`, or some other source. See
`RunningCovariance.from_stats` for a convenient way to construct these.
Args:
inner_kernel: `TransitionKernel`-like object.
initial_running_variance:
`tfp.experimental.stats.RunningVariance`-like object, or list of them,
for a batch of momentum distributions. These use `update` on the state
to maintain an estimate of the variance, and so space, and so must have
a structure compatible with the state space.
num_estimation_steps: An optional scalar `int` `Tensor` number of initial
steps to during which to adjust the running variance. This may be
greater, less than, or equal to the number of burnin steps. If this
argument is None, the mass matrix will be updated at each `one_step`
call. Otherwise, the mass matrix will be updated when the current step
is equal to `num_estimation_steps`.
momentum_distribution_setter_fn: A callable with the signature
`(kernel_results, new_momentum_distribution) -> new_kernel_results`
where `kernel_results` are the results of the `inner_kernel`,
`new_momentum_distribution` is a `CompositeTensor` or a nested
collection of `CompositeTensor`s, and `new_kernel_results` are a
possibly-modified copy of `kernel_results`. The default,
`hmc_like_momentum_distribution_setter_fn`, presumes HMC-style
`kernel_results`, and sets the `momentum_distribution` only under the
`accepted_results` field.
momentum_distribution_getter_fn: A callable with the signature
`kernel_results -> momentum_distribution`
where `kernel_results` are the results of the `inner_kernel` and
`momentum_distribution` is a `CompositeTensor` or a nested
collection of `CompositeTensor`s. The default,
`hmc_like_momentum_distribution_getter_fn`, presumes HMC-style
`kernel_results`, and gets the `momentum_distribution` only under the
`accepted_results` field.
validate_args: Python `bool`. When `True` kernel parameters are checked
for validity. When `False` invalid inputs may silently render incorrect
outputs.
experimental_shard_axis_names: An optional structure of string names
indicating how members of the state are sharded.
name: Python `str` name prefixed to Ops created by this class. Default:
'diagonal_mass_matrix_adaptation'.
"""
inner_kernel = mcmc_util.enable_store_parameters_in_results(
inner_kernel).experimental_with_shard_axes(
experimental_shard_axis_names)
self._parameters = dict(
inner_kernel=inner_kernel,
initial_running_variance=initial_running_variance,
num_estimation_steps=num_estimation_steps,
momentum_distribution_setter_fn=momentum_distribution_setter_fn,
momentum_distribution_getter_fn=momentum_distribution_getter_fn,
experimental_shard_axis_names=experimental_shard_axis_names,
name=name,
)
