View source on GitHub |
Adapts the inner kernel's momentum_distribution
to estimated variance.
Inherits From: TransitionKernel
tfp.experimental.mcmc.DiagonalMassMatrixAdaptation(
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
)
This kernel uses an online variance estimate to adjust a diagonal covariance
matrix for each of the state parts. More specifically, the
momentum_distribution
of the innermost kernel is set to a diagonal
multivariate normal distribution whose variance is the inverse of the
online estimate. The inverse of the covariance of the momentum is often called
the "mass matrix" in the context of Hamiltonian Monte Carlo.
This preconditioning scheme works well when the covariance is diagonally dominant, and may give reasonable results even when the number of draws is less than the dimension. In particular, it should generally do a better job than no preconditioning, which implicitly uses an identity mass matrix.
Note that this kernel does not implement a calibrated sampler; rather, it is intended to be used as one step of an iterative adaptation process. It should not be used when drawing actual samples.
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'.
|
Attributes | |
---|---|
experimental_shard_axis_names
|
The shard axis names for members of the state. |
initial_running_variance
|
|
inner_kernel
|
|
is_calibrated
|
Returns True if Markov chain converges to specified distribution.
|
name
|
|
num_estimation_steps
|
|
parameters
|
Return dict of __init__ arguments and their values.
|
Methods
bootstrap_results
bootstrap_results(
init_state
)
Returns an object with the same type as returned by one_step(...)[1]
.
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.
|
copy
copy(
**override_parameter_kwargs
)
Non-destructively creates a deep copy of the kernel.
Args | |
---|---|
**override_parameter_kwargs
|
Python String/value dictionary of
initialization arguments to override with new values.
|
Returns | |
---|---|
new_kernel
|
TransitionKernel object of same type as self ,
initialized with the union of self.parameters and
override_parameter_kwargs, with any shared keys overridden by the
value of override_parameter_kwargs, i.e.,
dict(self.parameters, **override_parameters_kwargs) .
|
experimental_with_shard_axes
experimental_with_shard_axes(
shard_axis_names
)
Returns a copy of the kernel with the provided shard axis names.
Args | |
---|---|
shard_axis_names
|
a structure of strings indicating the shard axis names for each component of this kernel's state. |
Returns | |
---|---|
A copy of the current kernel with the shard axis information. |
momentum_distribution_getter_fn
momentum_distribution_getter_fn(
kernel_results
)
momentum_distribution_setter_fn
momentum_distribution_setter_fn(
kernel_results, new_momentum_distribution
)
one_step
one_step(
current_state, previous_kernel_results, seed=None
)
Takes one step of the TransitionKernel.
Must be overridden by subclasses.
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.
|