def testBatchedLasVegasAlgorithm(self):
seed = test_util.test_seed()
def uniform_less_than_point_five(seed):
seed_stream = SeedStream(seed, 'uniform_less_than_point_five')
values = tf.random.uniform([6], seed=seed_stream())
negative_values = -values
good = tf.less(values, 0.5)
return ((negative_values, values), good)
((negative_values, values), _) = self.evaluate(
brs.batched_las_vegas_algorithm(uniform_less_than_point_five,
seed=seed))
self.assertAllLess(values, 0.5)
self.assertAllClose(-values, negative_values)
if tf.executing_eagerly():
tf.random.set_seed(seed)
# Check for reproducibility.
((negative_values_2, values_2), _) = self.evaluate(
brs.batched_las_vegas_algorithm(uniform_less_than_point_five,
seed=seed))
self.assertAllEqual(negative_values, negative_values_2)
self.assertAllEqual(values, values_2)
def testGivingUp(self):
def trial(seed):
del seed
return tf.constant([1, 0]), tf.constant([True, False])
values, final_successes, num_iters = self.evaluate(
brs.batched_las_vegas_algorithm(trial,
max_trials=50,
seed=test_util.test_seed()))
self.assertAllEqual(values, [1, 0])
self.assertAllEqual(final_successes, [True, False])
self.assertAllEqual(50, num_iters)
def testBatchedLasVegasAlgorithm(self):
def uniform_less_than_point_five(seed):
values = samplers.uniform([6], seed=seed)
negative_values = -values
good = tf.less(values, 0.5)
return ((negative_values, values), good)
((negative_values, values), _, _) = self.evaluate(
brs.batched_las_vegas_algorithm(uniform_less_than_point_five,
seed=test_util.test_seed()))
self.assertAllLess(values, 0.5)
self.assertAllClose(-values, negative_values)
# Check for reproducibility.
((negative_values_2, values_2), _, _) = self.evaluate(
brs.batched_las_vegas_algorithm(uniform_less_than_point_five,
seed=test_util.test_seed()))
self.assertAllEqual(negative_values, negative_values_2)
self.assertAllEqual(values, values_2)
def _btrs(counts, probs, full_shape, seed):
"""Binomial transformed rejection sampler, for count*prob >= 10."""
# We use a transformation-rejection algorithm from
# pairs of uniform random variables due to Hormann.
# https://www.tandfonline.com/doi/abs/10.1080/00949659308811496
seed = samplers.sanitize_seed(seed)
# This is spq in the paper.
stddev = tf.math.sqrt(counts * probs * (1 - probs))
# Other coefficients for Transformed Rejection sampling.
b = 1.15 + 2.53 * stddev
a = -0.0873 + 0.0248 * b + 0.01 * probs
c = counts * probs + 0.5
r = probs / (1 - probs)
alpha = (2.83 + 5.1 / b) * stddev
m = tf.math.floor((counts + 1) * probs)
def batched_las_vegas_trial_fn(seed):
u_seed, v_seed = samplers.split_seed(seed)
u = samplers.uniform(full_shape, seed=u_seed, dtype=counts.dtype) - 0.5
v = samplers.uniform(full_shape, seed=v_seed, dtype=counts.dtype)
us = 0.5 - tf.math.abs(u)
k = tf.math.floor((2 * a / us + b) * u + c)
# When the bounding box is tight, this criteria is more numerically stable
# and equally valid. Particularly on GPU/TPU, it may make the difference
# between terminating and non-terminating loops.
v_r = 0.92 - 4.2 / b
accept_boxed = (us >= 0.07) & (v <= v_r)
# Reject non-sensical answers.
reject = (k < 0) | (k > counts)
# This deviates from Hormann's BTRS algorithm, as there is a log missing.
# For all (u, v) pairs outside of the bounding box, this calculates the
# transformed-reject ratio.
v = tf.math.log(v * alpha / (a / (us * us) + b))
upperbound = ((m + 0.5) * tf.math.log(
(m + 1) / (r * (counts - m + 1))) + (counts + 1) * tf.math.log(
(counts - m + 1) / (counts - k + 1)) +
(k + 0.5) * tf.math.log(r * (counts - k + 1) / (k + 1)) +
_stirling_approx_tail(m) +
_stirling_approx_tail(counts - m) -
_stirling_approx_tail(k) -
_stirling_approx_tail(counts - k))
accept_bounded = v <= upperbound
return k, (~reject) & (accept_boxed | accept_bounded)
return batched_rejection_sampler.batched_las_vegas_algorithm(
batched_las_vegas_trial_fn, seed=seed)[0] # Pick out samples.
def _btrs(counts, probs, full_shape, seed):
"""Binomial transformed rejection sampler, for count*prob >= 10."""
# We use a transformation-rejection algorithm from
# pairs of uniform random variables due to Hormann.
# https://www.tandfonline.com/doi/abs/10.1080/00949659308811496
seed = samplers.sanitize_seed(seed)
# This is spq in the paper.
stddev = tf.math.sqrt(counts * probs * (1 - probs))
# Other coefficients for Transformed Rejection sampling.
b = 1.15 + 2.53 * stddev
a = -0.0873 + 0.0248 * b + 0.01 * probs
c = counts * probs + 0.5
r = probs / (1 - probs)
alpha = (2.83 + 5.1 / b) * stddev
m = tf.math.floor((counts + 1) * probs)
def batched_las_vegas_trial_fn(seed):
u_seed, v_seed = samplers.split_seed(seed)
u = samplers.uniform(full_shape, seed=u_seed, dtype=counts.dtype) - 0.5
v = samplers.uniform(full_shape, seed=v_seed, dtype=counts.dtype)
us = 0.5 - tf.math.abs(u)
k = tf.math.floor((2 * a / us + b) * u + c)
# The original algorithm accepts early if (us >= 0.07) & (v <= v_r), where
# `v_r = 0.92 - 4.2 / b`, the region for which the box is tight. Since we
# have rewritten from scalar evaluation to batch and we can't avoid the
# computation below, we omit the cheaper check.
# Reject non-sensical answers.
reject = (k < 0) | (k > counts)
# This deviates from Hormann's BTRS algorithm, as there is a log missing.
# For all (u, v) pairs outside of the bounding box, this calculates the
# transformed-reject ratio.
v = tf.math.log(v * alpha / (a / (us * us) + b))
upperbound = ((m + 0.5) * tf.math.log(
(m + 1) / (r * (counts - m + 1))) + (counts + 1) * tf.math.log(
(counts - m + 1) / (counts - k + 1)) +
(k + 0.5) * tf.math.log(r * (counts - k + 1) / (k + 1)) +
_stirling_approx_tail(m) +
_stirling_approx_tail(counts - m) -
_stirling_approx_tail(k) -
_stirling_approx_tail(counts - k))
return k, (~reject) & (v <= upperbound)
return batched_rejection_sampler.batched_las_vegas_algorithm(
batched_las_vegas_trial_fn, seed)[0] # Drop `num_trials`.
def generate_positive_v():
"""Generate positive v."""
def _inner(seed):
x = samplers.normal(shape, dtype=internal_dtype, seed=seed)
# This implicitly broadcasts concentration up to sample shape.
v = 1 + c * x
return (x, v), v > 0.
# Note: It should be possible to remove this 'inner' call to
# `batched_las_vegas_algorithm` and merge the v > 0 check into the
# overall check for a good sample. This would lead to a slightly simpler
# implementation; it is unclear whether it would be faster. We include
# the inner loop so this implementation is more easily comparable to
# Ref. [1] and other implementations.
return brs.batched_las_vegas_algorithm(_inner, v_seed)[0]
def testBatchedStructuredLasVegasAlgorithm(self):
def uniform_in_circle(seed):
coords = samplers.uniform([6, 2],
minval=-1.0,
maxval=1.0,
seed=seed)
radii = tf.reduce_sum(coords * coords, axis=-1)
good = tf.less(radii, 1)
return (coords, good)
(coords, _,
_) = brs.batched_las_vegas_algorithm(uniform_in_circle,
seed=test_util.test_seed())
radii = self.evaluate(tf.reduce_sum(coords * coords, axis=-1))
self.assertAllLess(radii, 1.0)
def _random_poisson_high_rate(sample_shape,
log_rate,
internal_dtype=tf.float64,
seed=None):
"""Samples from the Poisson distribution using transformed rejection sampling.
Given a CDF F(x), and G(x), a dominating distribution chosen such that it is
close to the inverse CDF F^-1(x), compute the following steps:
1) Generate U and V, two independent random variates. Set U = U - 0.5 (this
step isn't strictly necessary, but is done to make some calculations symmetric
and convenient. Henceforth, G is defined on [-0.5, 0.5]).
2) If V <= alpha * F'(G(U)) * G'(U), return floor(G(U)), else return to
step 1. alpha is the acceptance probability of the rejection algorithm.
The dominating distribution in this case:
G(u) = (2 * a / (2 - |u|) + b) * u + c
For more details on transformed rejection, see [1].
Args:
sample_shape: The output sample shape. Must broadcast with `log_rate`.
log_rate: Floating point tensor, log rate.
internal_dtype: dtype to use for internal computations.
seed: (optional) The random seed.
Returns:
Samples from the poisson distribution using transformed rejection.
#### References
[1]: W. Hormann, G. Derflinger, The Transformed Rejection Method For
Generating Random Variables, An Alternative To The Ratio Of Uniforms Method
(1994), Manuskript, Institut f. Statistik, Wirtschaftsuniversitat
"""
rate = tf.math.exp(log_rate)
b = 0.931 + 2.53 * tf.math.exp(0.5 * log_rate)
a = -0.059 + 0.02483 * b
inverse_alpha = 1.1239 + 1.1328 / (b - 3.4)
def generate_and_test_samples(seed):
"""Generate and test samples."""
u_seed, v_seed = samplers.split_seed(seed)
u = samplers.uniform(sample_shape, dtype=internal_dtype, seed=u_seed)
u = u - 0.5
u_shifted = 0.5 - tf.math.abs(u)
v = samplers.uniform(sample_shape, dtype=internal_dtype, seed=v_seed)
k = tf.math.floor(((2. * a) / u_shifted + b) * u + rate + 0.43)
good_sample_mask = (u_shifted >= 0.07) & (v <= 0.9277 - 3.6224 / (b - 2.))
s = tf.math.log(v * inverse_alpha / (a / tf.math.square(u_shifted) + b))
t = -rate + k * log_rate - tf.math.lgamma(k + 1)
good_sample_mask = good_sample_mask | (s <= t)
# Make sure the sample is within bounds.
good_sample_mask = good_sample_mask & (k >= 0) & ((u_shifted >= 0.013) |
(v <= u_shifted))
return k, good_sample_mask
samples = brs.batched_las_vegas_algorithm(
generate_and_test_samples, seed=seed)[0]
return samples
def rejection_sample(concentration):
"""Gamma rejection sampler."""
# Note, concentration here already has a shape that is broadcast with rate.
cast_concentration = tf.cast(concentration, internal_dtype)
good_params_mask = (concentration >= 0.)
# When replacing NaN values, use 100. for concentration, since that leads to
# a high-likelihood of the rejection sampler accepting on the first pass.
safe_concentration = tf.where(good_params_mask, cast_concentration,
100.)
modified_safe_concentration = tf.where(safe_concentration < 1.,
safe_concentration + 1.,
safe_concentration)
one_third = tf.constant(1. / 3, dtype=internal_dtype)
d = modified_safe_concentration - one_third
c = one_third * tf.math.rsqrt(d)
def generate_and_test_samples(seed):
"""Generate and test samples."""
v_seed, u_seed = samplers.split_seed(seed)
x = samplers.normal(shape, dtype=internal_dtype, seed=v_seed)
# This implicitly broadcasts concentration up to sample shape.
v = 1 + c * x
# In [1], there is an 'inner' rejection sampling loop which checks that
# v > 0 and generates a new normal sample if it's not, saving the rest of
# the computations below. We found that merging the check for v > 0 with
# the `good_sample_mask` not only simplifies the code, but leads to a
# ~2x speedup for small concentrations on GPU, at the cost of deviating
# slightly from the implementation given in Ref. [1].
accept_v = v > 0.
logv = tf.math.log1p(c * x)
x2 = x * x
v3 = v * v * v
logv3 = logv * 3
u = samplers.uniform(shape, dtype=internal_dtype, seed=u_seed)
# In [1], the suggestion is to first check u < 1 - 0.331 * x2 * x2, and to
# run the check below only if it fails, in order to avoid the relatively
# expensive logarithm calls. Our algorithm operates in batch mode: we will
# have to compute or not compute the logarithms for the entire batch, and
# as the batch gets larger, the odds we compute it grow. Therefore we
# don't bother with the "cheap" check.
good_sample_mask = tf.logical_and(
tf.math.log(u) < (x2 / 2. + d * (1 - v3 + logv3)), accept_v)
return logv3 if log_space else v3, good_sample_mask
samples = brs.batched_las_vegas_algorithm(
generate_and_test_samples, seed=generate_and_test_samples_seed)[0]
concentration_fix_unif = samplers.uniform( # in [0, 1)
shape,
dtype=internal_dtype,
seed=concentration_fix_seed)
if log_space:
concentration_lt_one_fix = tf.where(
safe_concentration < 1.,
# Why do we use log1p(-x)? x is in [0, 1) and log(0) = -inf, is bad.
# x ~ U(0,1) => 1-x ~ U(0,1)
# But at the boundary, 1-x in (0, 1]. Good.
# So we can take log(unif(0,1)) safely as log(1-unif(0,1)).
# log1p(-0) = 0, and log1p(-almost_one) = -not_quite_inf. Good.
tf.math.log1p(-concentration_fix_unif) / safe_concentration,
tf.zeros((), dtype=internal_dtype))
samples = samples + tf.math.log(d) + concentration_lt_one_fix
else:
concentration_lt_one_fix = tf.where(
safe_concentration < 1.,
tf.math.pow(concentration_fix_unif,
tf.math.reciprocal(safe_concentration)),
tf.ones((), dtype=internal_dtype))
samples = samples * d * concentration_lt_one_fix
samples = tf.where(good_params_mask, samples, np.nan)
output_type_samples = tf.cast(samples, output_dtype)
return output_type_samples
def rejection_sample(concentration):
"""Gamma rejection sampler."""
# Note, concentration here already has a shape that is broadcast with rate.
cast_concentration = tf.cast(concentration, internal_dtype)
good_params_mask = (concentration > 0.)
# When replacing NaN values, use 100. for concentration, since that leads to
# a high-likelihood of the rejection sampler accepting on the first pass.
safe_concentration = tf.where(good_params_mask, cast_concentration, 100.)
modified_safe_concentration = tf.where(
safe_concentration < 1., safe_concentration + 1., safe_concentration)
one_third = tf.constant(1. / 3, dtype=internal_dtype)
d = modified_safe_concentration - one_third
c = one_third * tf.math.rsqrt(d)
def generate_and_test_samples(seed):
"""Generate and test samples."""
v_seed, u_seed = samplers.split_seed(seed)
def generate_positive_v():
"""Generate positive v."""
def _inner(seed):
x = samplers.normal(shape, dtype=internal_dtype, seed=seed)
# This implicitly broadcasts concentration up to sample shape.
v = 1 + c * x
return (x, v), v > 0.
# Note: It should be possible to remove this 'inner' call to
# `batched_las_vegas_algorithm` and merge the v > 0 check into the
# overall check for a good sample. This would lead to a slightly simpler
# implementation; it is unclear whether it would be faster. We include
# the inner loop so this implementation is more easily comparable to
# Ref. [1] and other implementations.
return brs.batched_las_vegas_algorithm(_inner, v_seed)[0]
(x, v) = generate_positive_v()
logv = tf.math.log1p(c * x)
x2 = x * x
v3 = v * v * v
logv3 = logv * 3
u = samplers.uniform(
shape, dtype=internal_dtype, seed=u_seed)
# In [1], the suggestion is to first check u < 1 - 0.331 * x2 * x2, and to
# run the check below only if it fails, in order to avoid the relatively
# expensive logarithm calls. Our algorithm operates in batch mode: we will
# have to compute or not compute the logarithms for the entire batch, and
# as the batch gets larger, the odds we compute it grow. Therefore we
# don't bother with the "cheap" check.
good_sample_mask = tf.math.log(u) < (x2 / 2. + d * (1 - v3 + logv3))
return logv3 if log_space else v3, good_sample_mask
samples = brs.batched_las_vegas_algorithm(
generate_and_test_samples, seed=generate_and_test_samples_seed)[0]
concentration_fix_unif = samplers.uniform( # in [0, 1)
shape, dtype=internal_dtype, seed=concentration_fix_seed)
if log_space:
concentration_lt_one_fix = tf.where(
safe_concentration < 1.,
# Why do we use log1p(-x)? x is in [0, 1) and log(0) = -inf, is bad.
# x ~ U(0,1) => 1-x ~ U(0,1)
# But at the boundary, 1-x in (0, 1]. Good.
# So we can take log(unif(0,1)) safely as log(1-unif(0,1)).
# log1p(-0) = 0, and log1p(-almost_one) = -not_quite_inf. Good.
tf.math.log1p(-concentration_fix_unif) / safe_concentration,
tf.zeros((), dtype=internal_dtype))
samples = samples + tf.math.log(d) + concentration_lt_one_fix
else:
concentration_lt_one_fix = tf.where(
safe_concentration < 1.,
tf.math.pow(concentration_fix_unif,
tf.math.reciprocal(safe_concentration)),
tf.ones((), dtype=internal_dtype))
samples = samples * d * concentration_lt_one_fix
samples = tf.where(good_params_mask, samples, np.nan)
output_type_samples = tf.cast(samples, output_dtype)
return output_type_samples
def rejection_sample(alpha):
"""Gamma rejection sampler."""
# Note that alpha here already has a shape that is broadcast with beta.
cast_alpha = tf.cast(alpha, internal_dtype)
good_params_mask = (alpha > 0.)
# When replacing NaN values, use 100. for alpha, since that leads to a
# high-likelihood of the rejection sampler accepting on the first pass.
safe_alpha = tf.where(good_params_mask, cast_alpha, 100.)
modified_safe_alpha = tf.where(safe_alpha < 1., safe_alpha + 1.,
safe_alpha)
one_third = tf.constant(1. / 3, dtype=internal_dtype)
d = modified_safe_alpha - one_third
c = one_third / tf.sqrt(d)
def generate_and_test_samples(seed):
"""Generate and test samples."""
seed_stream = SeedStream(seed, 'generate_and_test_samples')
def generate_positive_v():
"""Generate positive v."""
def _inner(seed):
seed_stream = SeedStream(seed, '_inner')
x = tf.random.normal(sample_shape,
dtype=internal_dtype,
seed=seed_stream())
# This implicitly broadcasts alpha up to sample shape.
v = 1 + c * x
return (x, v), v > 0.
# Note: It should be possible to remove this 'inner' call to
# `batched_las_vegas_algorithm` and merge the v > 0 check into the
# overall check for a good sample. This would lead to a slightly simpler
# implementation; it is unclear whether it would be faster. We include
# the inner loop so this implementation is more easily comparable to
# Ref. [1] and other implementations.
return brs.batched_las_vegas_algorithm(_inner, seed)[0]
(x, v) = generate_positive_v()
x2 = x * x
v3 = v * v * v
u = tf.random.uniform(sample_shape,
dtype=internal_dtype,
seed=seed_stream())
# In [1], the suggestion is to first check u < 1 - 0.331 * x2 * x2, and to
# run the check below only if it fails, in order to avoid the relatively
# expensive logarithm calls. Our algorithm operates in batch mode: we will
# have to compute or not compute the logarithms for the entire batch, and
# as the batch gets larger, the odds we compute it grow. Therefore we
# don't bother with the "cheap" check.
good_sample_mask = (tf.math.log(u) < x2 / 2. + d *
(1 - v3 + tf.math.log(v3)))
return v3, good_sample_mask
samples = brs.batched_las_vegas_algorithm(generate_and_test_samples,
seed=seed_stream())[0]
samples = samples * d
one = tf.constant(1., dtype=internal_dtype)
alpha_lt_one_fix = tf.where(
safe_alpha < 1.,
tf.math.pow(
tf.random.uniform(sample_shape,
dtype=internal_dtype,
seed=seed_stream()), one / safe_alpha), one)
samples = samples * alpha_lt_one_fix
samples = tf.where(good_params_mask, samples, np.nan)
output_type_samples = tf.cast(samples, output_dtype)
# We use `tf.where` instead of `tf.maximum` because we need to allow for
# `samples` to be `nan`, but `tf.maximum(nan, x) == x`.
output_type_samples = tf.where(
output_type_samples == 0,
np.finfo(dtype_util.as_numpy_dtype(
output_type_samples.dtype)).tiny, output_type_samples)
def grad(dy):
"""The gradient of the normalized (beta=1) gamma samples w.r.t alpha."""
# Recall that cast_alpha has shape broadcast with beta, and samples and dy
# have shape sample_shape (which further expands the alpha-beta broadcast
# shape on the left).
cast_dy = tf.cast(dy, internal_dtype)
partial_alpha = tf.raw_ops.RandomGammaGrad(alpha=cast_alpha,
sample=samples)
grad = tf.cast(
tf.math.reduce_sum(
cast_dy * partial_alpha,
axis=tf.range(tf.rank(partial_alpha) - tf.rank(alpha))),
output_dtype)
return grad
return output_type_samples, grad # rejection_sample
def retry_init(proposal_fn, target_fn, *args, max_trials=50,
seed=None, name=None, **kwargs):
"""Tries an MCMC initialization proposal until it gets a valid state.
In this case, "valid" is defined as the value of `target_fn` is
finite. This corresponds to an MCMC workflow where `target_fn`
compute the log-probability one wants to sample from, in which case
"finite `target_fn`" means "finite and positive probability state".
If `target_fn` returns a Tensor of size greater than 1, the results
are assumed to be independent of each other, so that different batch
members can be accepted individually.
The method is bounded rejection sampling. The bound serves to avoid
wasting computation on hopeless initialization procedures. In
interactive MCMC, one would presumably rather come up with a better
initialization proposal than wait for an unbounded number of
attempts with a bad one. If unbounded re-trials are desired,
set `max_trials` to `None`.
Note: XLA and @jax.jit do not support assertions, so this function
can return invalid states on those platforms without raising an
error (unless `max_trials` is set to `None`).
Args:
proposal_fn: A function accepting a `seed` keyword argument and no other
required arguments which generates proposed initial states.
target_fn: A function accepting the return value of `proposal_fn`
and returning a floating-point Tensor.
*args: Additional arguments passed to `proposal_fn`.
max_trials: Size-1 integer `Tensor` or None. Maximum number of
calls to `proposal_fn` to attempt. If acceptable states are not
found in this many trials, `retry_init` signals an error. If
`None`, there is no limit, and `retry_init` skips the control
flow cost of checking for success.
seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
name: Python `str` name prefixed to Ops created by this function.
Default value: `None` (i.e., 'mcmc_sample_chain').
**kwargs: Additional keyword arguments passed to `proposal_fn`.
Returns:
states: An acceptable result from `proposal_fn`.
#### Example
One popular MCMC initialization scheme is to start the chains near 0
in unconstrained space. There are models where the unconstraining
transformation cannot exactly capture the space of valid states,
such that this initialization has some material but not overwhelming
chance of failure. In this case, we can use `retry_init` to compensate.
```python
@tfp.distributions.JointDistributionCoroutine
def model():
...
raw_init_dist = tfp.experimental.mcmc.init_near_unconstrained_zero(model)
init_states = tfp.experimental.mcmc.retry_init(
proposal_fn=raw_init_dist.sample,
target_fn=model.log_prob,
sample_shape=[100],
seed=[4, 8])
states = tfp.mcmc.sample_chain(
current_state=init_states,
...)
```
"""
def trial(seed):
values = proposal_fn(*args, seed=seed, **kwargs)
log_probs = target_fn(values)
success = tf.math.is_finite(log_probs)
return values, success
with tf.name_scope(name or 'mcmc_retry_init'):
values, successes, _ = brs.batched_las_vegas_algorithm(
trial, max_trials=max_trials, seed=seed)
if max_trials is None:
# We were authorized to compute until success, so no need to
# check for failure
return values
else:
num_states = tf.size(successes)
num_successes = tf.reduce_sum(tf.cast(successes, tf.int32))
msg = ('Failed to find acceptable initial states after {} trials;\n'
'{} of {} states have non-finite log probability').format(
max_trials, num_states - num_successes, num_states)
with tf.control_dependencies([tf.debugging.assert_equal(
successes, tf.ones_like(successes), message=msg)]):
return tf.nest.map_structure(tf.identity, values)
