View source on GitHub |
Convenience constructor for common MCMC transition kernels.
tfp.experimental.mcmc.KernelBuilder(
target_log_prob_fn,
core_class,
core_params,
default_step_size_on,
step_size,
step_adapter_class,
step_adapter_params,
default_target_accept_prob_on,
target_accept_prob,
replica_exchange_params,
transform_params,
num_steps_between_results,
auto_tracing_on,
tracing_params,
show_progress,
user_reducer
)
KernelBuilder
gives an alternative interface for building MCMC transition
kernels. It wraps the base tfp.mcmc
library, offering more convenience at
the cost of some power and flexibility.
It is designed to work in conjunction with KernelOutputs
for more
convenience.
Example usage:
# Initialize builder with `target_log_prob_fn`.
builder = KernelBuilder(target_log_prob_fn)
# Configure initial transition kernel.
builder = (
builder
# Use Hamilton Monte Carlo
.hmc(num_leapfrog_steps=3)
# with step size adaptation
.dual_averaging_adaptation()
.set_num_adaptation_steps(50)
# and a transformation
.transform(my_bijector))
# Do sampling...
outputs = builder.sample(num_steps, initial_state)
# Start from the previous `KernelBuilder` configuration.
builder = (
builder
# Continue using HMC...
# Still use `my_bijector` transform
# But with no step size adaptation:
.clear_step_adapter()
# set a static step size
.set_step_size(outputs.new_step_size))
# More sampling starting from where we left off.
outputs2 = builder.sample(num_steps, outputs.current_state)
# Etc ...
All methods except build()
and sample()
return a modified copy of the
builder for further method-chaining. The builder itself is immutable
(namedtuple) for safe use inside of graphs.
KernelBuilder
builds kernels with the following parts (in order):
- A core transition kernel
- Optional step size adaptation or replica exchange
- Transformating bijector
- Thinning
- Streaming reductions
The core kernels can be either HamiltonianMonteCarlo
, NoUTurnSampler
,
PreconditionedHamiltonianMonteCarlo
, MetropolisAdjustedLangevinAlgorithm
,
or RandomWalkMetropolis
. Support for additional core kernels may be added
in the future.
Step size adaption is performed by SimpleStepSizeAdaptation
or
DualAveragingStepSizeAdaptation
. Note not all core kernels are currently
compatible with step size adaptation.
KernelBuilder
maintains some state between kernel builds which can be reused
or overriden:
- Target log prob function
- Core kernel class
- Step size (initial)
- Step size adapter class (optional) 4a. Number of adaptation steps 4b. Target acceptance probability
- Replica exchange parameters (optional)
TransformedTransitionKernel
bijector/params (optional)- Thinning: number of steps between results (optional)
- Tracing parameters for
TracingReducer
/ auto-tracing. - Show progress boolean.
- Reductions for
WithReductions
See instance method documentation for more information.
Methods
build
build(
num_steps=None
)
Build and return the specified kernel.
Args | |
---|---|
num_steps
|
An integer. Some kernel pieces (step adaptation) require knowing the number of steps to sample in advance; pass that in here. |
Returns | |
---|---|
kernel
|
The configured TransitionKernel .
|
clear_reducer
clear_reducer()
Remove previously set reductions.
clear_replica_exchange
clear_replica_exchange()
clear_step_adapter
clear_step_adapter()
Removes step adaptation.
clear_tracing
clear_tracing()
Remove TracingReducer
.
clear_transform
clear_transform()
Remove previously set TransformedTransitionKernel
.
dual_averaging_adaptation
dual_averaging_adaptation(
exploration_shrinkage=0.05,
shrinkage_target=None,
step_count_smoothing=10,
decay_rate=0.75,
step_size_setter_fn=simple_step_size_adaptation.hmc_like_step_size_setter_fn,
step_size_getter_fn=simple_step_size_adaptation.hmc_like_step_size_getter_fn,
log_accept_prob_getter_fn=simple_step_size_adaptation.hmc_like_log_accept_prob_getter_fn,
validate_args=False,
name=None
)
Use DualAveragingStepSizeAdaptation
.
See the DualAveragingStepSizeAdaptation
docs for more details.
Args | |
---|---|
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.
|
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').
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
get_step_size
get_step_size()
Return the set or default step size.
get_target_accept_prob
get_target_accept_prob()
Return the set target_accept_prob
or the default for the core kernel.
hmc
hmc(
num_leapfrog_steps,
state_gradients_are_stopped=False,
store_parameters_in_results=False,
name=None
)
Use the HamiltonianMonteCarlo
core transition kernel.
See the HamiltonianMonteCarlo
docs for more details.
Args | |
---|---|
num_leapfrog_steps
|
Integer number of steps to run the leapfrog integrator
for. Total progress per HMC step is roughly proportional to
step_size * num_leapfrog_steps .
|
state_gradients_are_stopped
|
Python bool indicating that the proposed
new state be run through tf.stop_gradient . This is particularly useful
when combining optimization over samples from the HMC chain.
Default value: False (i.e., do not apply stop_gradient ).
|
store_parameters_in_results
|
If True , then step_size and
num_leapfrog_steps are written to and read from eponymous fields in
the kernel results objects returned from one_step and
bootstrap_results . This allows wrapper kernels to adjust those
parameters on the fly.
|
name
|
Python str name prefixed to Ops created by this function.
Default value: None (e.g., 'hmc_kernel').
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
make
@classmethod
make( target_log_prob_fn )
Construct a KernelBuilder
with empty defaults.
mala
mala(
volatility_fn=None, parallel_iterations=10, name=None
)
Use the MetropolisAdjustedLangevinAlgorithm
core transition kernel.
See the MetropolisAdjustedLangevinAlgorithm
docs for more details.
Args | |
---|---|
volatility_fn
|
Python callable which takes an argument like
current_state (or *current_state if it's a list) and returns
volatility value at current_state . Should return a Tensor or Python
list of Tensor s that must broadcast with the shape of
current_state Defaults to the identity function.
|
parallel_iterations
|
the number of coordinates for which the gradients of
the volatility matrix volatility_fn can be computed in parallel.
|
name
|
Python str name prefixed to Ops created by this function.
Default value: None (i.e., 'mala_kernel').
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
nuts
nuts(
max_tree_depth=10,
max_energy_diff=1000.0,
unrolled_leapfrog_steps=1,
parallel_iterations=10,
name=None
)
Use the NoUTurnSampler
core kernel.
See the NoUTurnSampler
docs for more details.
Args | |
---|---|
max_tree_depth
|
Maximum depth of the tree implicitly built by NUTS. The
maximum number of leapfrog steps is bounded by 2**max_tree_depth i.e.
the number of nodes in a binary tree max_tree_depth nodes deep. The
default setting of 10 takes up to 1024 leapfrog steps.
|
max_energy_diff
|
Scaler threshold of energy differences at each leapfrog, divergence samples are defined as leapfrog steps that exceed this threshold. Default to 1000. |
unrolled_leapfrog_steps
|
The number of leapfrogs to unroll per tree expansion step. Applies a direct linear multipler to the maximum trajectory length implied by max_tree_depth. Defaults to 1. |
parallel_iterations
|
The number of iterations allowed to run in parallel.
It must be a positive integer. See tf.while_loop for more details.
|
name
|
Python str name prefixed to Ops created by this function.
Default value: None (i.e., 'nuts_kernel').
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
phmc
phmc(
num_leapfrog_steps,
momentum_distribution=None,
state_gradients_are_stopped=False,
store_parameters_in_results=False,
name=None
)
Use the PreconditionedHamiltonianMonteCarlo
core transition kernel.
See the PreconditionedHamiltonianMonteCarlo
docs for more details.
Args | |
---|---|
num_leapfrog_steps
|
Integer number of steps to run the leapfrog integrator
for. Total progress per HMC step is roughly proportional to
step_size * num_leapfrog_steps .
|
momentum_distribution
|
A tfp.distributions.Distribution instance to draw
momentum from. Defaults to isotropic normal distributions.
|
state_gradients_are_stopped
|
Python bool indicating that the proposed
new state be run through tf.stop_gradient . This is particularly useful
when combining optimization over samples from the HMC chain.
Default value: False (i.e., do not apply stop_gradient ).
|
store_parameters_in_results
|
If True , then step_size ,
momentum_distribution , and num_leapfrog_steps are written to and
read from eponymous fields in the kernel results objects returned from
one_step and bootstrap_results . This allows wrapper kernels to
adjust those parameters on the fly. In case this is True , the
momentum_distribution must be a CompositeTensor . See
tfp.experimental.auto_composite . This is incompatible with
step_size_update_fn , which must be set to None .
|
name
|
Python str name prefixed to Ops created by this function.
Default value: None (i.e., 'hmc_kernel').
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
replica_exchange
replica_exchange(
inverse_temperatures,
swap_proposal_fn=replica_exchange_mc.default_swap_proposal_fn(1.0),
state_includes_replicas=False,
validate_args=False,
name=None
)
Use ReplicaExchangeMC
.
See the ReplicaExchangeMC
docs for more details.
Args | |
---|---|
inverse_temperatures
|
Tensor of inverse temperatures to temper each
replica. The leftmost dimension is the num_replica and the
second dimension through the rightmost can provide different temperature
to different batch members, doing a left-justified broadcast.
|
swap_proposal_fn
|
Python callable which take a number of replicas, and
returns swaps , a shape [num_replica] + batch_shape Tensor , where
axis 0 indexes a permutation of {0,..., num_replica-1} , designating
replicas to swap.
|
state_includes_replicas
|
Boolean indicating whether the leftmost dimension
of each state sample should index replicas. If True , the leftmost
dimension of the current_state kwarg to tfp.mcmc.sample_chain will
be interpreted as indexing replicas.
|
validate_args
|
Python bool , default False . When True distribution
parameters are checked for validity despite possibly degrading runtime
performance. 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., "remc_kernel").
|
Raises | |
---|---|
ValueError
|
inverse_temperatures doesn't have statically known 1D shape.
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
rwm
rwm(
new_state_fn=None, name=None
)
Use the RandomWalkMetropolis
core kernel.
See the RandomWalkMetropolis
docs for more details.
Args | |
---|---|
new_state_fn
|
Python callable which takes a list of state parts and a
seed; returns a same-type list of Tensor s, each being a perturbation
of the input state parts. The perturbation distribution is assumed to be
a symmetric distribution centered at the input state part.
Default value: None which is mapped to
tfp.mcmc.random_walk_normal_fn() .
|
name
|
Python str name prefixed to Ops created by this function.
Default value: None (i.e., 'rwm_kernel').
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
sample
sample(
num_steps, current_state, previous_kernel_results=None
)
Sample from the configured kernel.
Args | |
---|---|
num_steps
|
Integer number of Markov chain steps. |
current_state
|
Tensor or Python list of Tensor s representing the
current state(s) of the Markov chain(s).
|
previous_kernel_results
|
A Tensor or a nested collection of Tensor s.
Warm-start for the auxiliary state needed by the given kernel .
If not supplied, step_kernel will cold-start with
kernel.bootstrap_results .
|
Returns | |
---|---|
outputs
|
A KernelOutputs object containing the states, trace, etc.
|
set_auto_tracing
set_auto_tracing(
on=True
)
Add smart tracing.
set_num_steps_between_results
set_num_steps_between_results(
num_steps_between_results
)
Thin sampling by num_steps_between_results
.
set_reducer
set_reducer(
reducer
)
Use tfp.experimental.mcmc.WithReductions
.
See the WithReductions
docs for more details.
Args | |
---|---|
reducer
|
A (possibly nested) structure of Reducer s to be evaluated
on the inner_kernel 's samples.
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
set_show_progress
set_show_progress(
on=True
)
set_step_size
set_step_size(
step_size
)
Set the step size (for core kernels with a step size.)
set_target_accept_prob
set_target_accept_prob(
target_accept_prob
)
Set the target acceptance for step adaptation kernels.
Args | |
---|---|
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]. By default, this is
0.25 for RandomWalkMetropolis and 0.75 for HamiltonianMonteCarlo ,
MetropolisAdjustedLangevinAlgorithm and NoUTurnSampler .
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
set_tracing
set_tracing(
trace_fn=_trace_all_results, size=None, name=None
)
Trace sampling state and results.
See the TracingReducer
docs for more details.
Args | |
---|---|
trace_fn
|
A callable that takes in the current chain state and the
previous kernel results and return a Tensor or a nested collection
of Tensor s that is accumulated across samples.
|
size
|
Integer or scalar Tensor denoting the size of the accumulated
TensorArray . If this is None (which is the default), a
dynamic-shaped TensorArray will be used.
|
name
|
Python str name prefixed to Ops created by this function.
Default value: None (i.e., 'tracing_reducer').
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
simple_adaptation
simple_adaptation(
adaptation_rate=0.01,
step_size_setter_fn=simple_step_size_adaptation.hmc_like_step_size_setter_fn,
step_size_getter_fn=simple_step_size_adaptation.hmc_like_step_size_getter_fn,
log_accept_prob_getter_fn=simple_step_size_adaptation.hmc_like_log_accept_prob_getter_fn,
validate_args=False,
name=None
)
Use SimpleStepSizeAdaptation
.
See the SimpleStepSizeAdaptation
docs for more details.
Args | |
---|---|
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'.
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
transform
transform(
bijector, name=None
)
Use TransformedTransitionKernel
.
See TransformedTransitionKernel
docs for more details.
Args | |
---|---|
bijector
|
tfp.distributions.Bijector or list of
tfp.distributions.Bijector s. These bijectors use forward to map the
inner_kernel state space to the state expected by
inner_kernel.target_log_prob_fn .
|
name
|
Python str name prefixed to Ops created by this function.
Default value: None (i.e., "transformed_kernel").
|
Returns | |
---|---|
self
|
Returns the builder for more method chaining. |
use_default_step_size
use_default_step_size(
on=True
)
Use default step size (or not.)
use_default_target_accept_prob
use_default_target_accept_prob(
on=True
)
Use per-core class default target acceptance probability (or not.)