def adaptive_hamiltonian_monte_carlo_init(
state: 'fun_mc.TensorNest',
target_log_prob_fn: 'fun_mc.PotentialFn',
step_size: 'fun_mc.FloatTensor' = 1e-2,
initial_mean: 'fun_mc.FloatNest' = 0.,
initial_scale: 'fun_mc.FloatNest' = 1.,
scale_smoothing_steps: 'fun_mc.IntTensor' = 10,
) -> 'AdaptiveHamiltonianMonteCarloState':
"""Initializes `AdaptiveHamiltonianMonteCarloState`.
Args:
state: Initial state of the chain.
target_log_prob_fn: Target log prob fn.
step_size: Initial scalar step size.
initial_mean: Initial mean for computing the running variance estimate. Must
broadcast structurally and tensor-wise with state.
initial_scale: Initial scale for computing the running variance estimate.
Must broadcast structurally and tensor-wise with state.
scale_smoothing_steps: How much weight to assign to the `initial_mean` and
`initial_scale`. Increase this to stabilize early adaptation.
Returns:
adaptive_hmc_state: State of the `adaptive_hamiltonian_monte_carlo_step`
`TransitionOperator`.
"""
hmc_state = fun_mc.hamiltonian_monte_carlo_init(state, target_log_prob_fn)
dtype = util.flatten_tree(hmc_state.state)[0].dtype
chain_ndims = len(hmc_state.target_log_prob.shape)
running_var_state = fun_mc.running_variance_init(
shape=util.map_tree(lambda s: s.shape[chain_ndims:], hmc_state.state),
dtype=util.map_tree(lambda s: s.dtype, hmc_state.state),
)
initial_mean = fun_mc.maybe_broadcast_structure(initial_mean, state)
initial_scale = fun_mc.maybe_broadcast_structure(initial_scale, state)
# It's important to add some smoothing here, as initial updates can be very
# different than the stationary distribution.
# TODO(siege): Add a pseudo-update functionality to avoid fiddling with the
# internals here.
running_var_state = running_var_state._replace(
num_points=util.map_tree(
lambda p: ( # pylint: disable=g-long-lambda
int(np.prod(hmc_state.target_log_prob.shape)) * tf.cast(
scale_smoothing_steps, p.dtype)),
running_var_state.num_points),
mean=util.map_tree( # pylint: disable=g-long-lambda
lambda m, init_m: tf.ones_like(m) * init_m, running_var_state.mean,
initial_mean),
variance=util.map_tree( # pylint: disable=g-long-lambda
lambda v, init_s: tf.ones_like(v) * init_s**2,
running_var_state.variance, initial_scale),
)
ssa_state = step_size_adaptation_init(
tf.convert_to_tensor(step_size, dtype=dtype))
return AdaptiveHamiltonianMonteCarloState(
hmc_state=hmc_state,
running_var_state=running_var_state,
ssa_state=ssa_state,
step=tf.zeros([], tf.int32))
def stochastic_gradient_ascent_hmc_init(
state: 'fun_mc.State',
target_log_prob_fn: 'fun_mc.PotentialFn',
init_trajectory_length: 'fun_mc.FloatTensor',
trajectory_length_params_init_fn:
'Callable[[fun_mc.FloatTensor], Any]' = default_trajectory_length_init):
"""Initialize Stochastic Gradient Ascent HMC state.
Args:
state: Initial Markov Chain state.
target_log_prob_fn: Target log prob fn.
init_trajectory_length: Initial trajectory length. Passed to
`trajectory_length_params_init_fn`.
trajectory_length_params_init_fn: Initializer for the trajectory length
parameters.
Returns:
sga_hmc_state: New Stochastic Gradient Ascent HMC state.
"""
init_trajectory_length = tf.convert_to_tensor(init_trajectory_length)
init_trajectory_length_params = trajectory_length_params_init_fn(
init_trajectory_length)
return StochasticGradientAscentHMCState(
hmc_state=fun_mc.hamiltonian_monte_carlo_init(state, target_log_prob_fn),
step=tf.ones([], tf.int32),
trajectory_length_params_opt_state=fun_mc.adam_init(
init_trajectory_length_params),
trajectory_length_params_rmean_state=fun_mc.running_mean_init(
util.map_tree(lambda x: x.shape, init_trajectory_length_params),
util.map_tree(lambda x: x.dtype, init_trajectory_length_params),
)._replace(mean=init_trajectory_length_params),
)
def hmc_step(trajectory_length, axis_name=()):
@tfp.experimental.distribute.JointDistributionCoroutine
def model():
z = yield root(tfd.Normal(0., 1))
yield tfp.experimental.distribute.Sharded(
tfd.Sample(tfd.Normal(z, 1.), 8), axis_name)
@tfp.experimental.distribute.JointDistributionCoroutine
def momentum_dist():
yield root(tfd.Normal(0., 2))
yield root(
tfp.experimental.distribute.Sharded(
tfd.Sample(tfd.Normal(0., 3.), 8), axis_name))
def target_log_prob_fn(x):
return model.log_prob(x), ()
def kinetic_energy_fn(m):
return -momentum_dist.log_prob(m), ()
def momentum_sample_fn(seed):
return momentum_dist.sample(2, seed=seed)
state = model.sample(2, seed=seed)
hmc_state = fun_mc.hamiltonian_monte_carlo_init(
state, target_log_prob_fn)
hmc_state, hmc_extra = (
prefab.hamiltonian_monte_carlo_with_state_grads_step(
hmc_state,
trajectory_length=trajectory_length,
scalar_step_size=epsilon,
step_size_scale=util.map_tree(lambda x: 1. + tf.abs(x),
state),
target_log_prob_fn=target_log_prob_fn,
seed=seed,
kinetic_energy_fn=kinetic_energy_fn,
momentum_sample_fn=momentum_sample_fn,
shard_axis_names=model.experimental_shard_axis_names))
def sum_state(x, axis_name):
res = tf.reduce_sum(x**2)
if axis_name:
res = backend.distribute_lib.psum(res, axis_name)
return res
sum_sq = util.map_tree_up_to(hmc_extra.proposed_state, sum_state,
hmc_extra.proposed_state,
model.experimental_shard_axis_names)
sum_sq = sum(util.flatten_tree(sum_sq))
return sum_sq, ()