def testBatchedRejectionBetaSample(self):
seed = test_util.test_seed()
# We build a rejection sampler for two beta distributions (in a batch): a
# beta(2, 5) and a beta(2, 2). Eyeballing an image on the wikipedia page,
# these are upper bounded by rectangles of heights 2.5 and 1.6 respectively.
alpha = np.array([2.], dtype=np.float32)
beta = np.array([5., 2.], dtype=np.float32)
upper_bounds = tf.constant([2.5, 1.6], dtype=tf.float32)
samples_per_distribution = 10000
target = tfd.Beta(alpha, beta).prob
def proposal(seed_stream):
seed_stream = SeedStream(seed, 'proposal')
uniform_samples = tfd.Uniform().sample(
[samples_per_distribution, 2], seed=seed_stream())
return uniform_samples, tf.ones_like(
uniform_samples) * upper_bounds
all_samples, _ = self.evaluate(
brs.batched_rejection_sampler(proposal, target, seed=seed))
for i in range(beta.shape[0]):
samples = all_samples[:, i]
ks, _ = sp_stats.kstest(samples, sp_stats.beta(alpha, beta[i]).cdf)
self.assertLess(ks, 0.02)
if tf.executing_eagerly():
tf.random.set_seed(seed)
# Check for reproducibility.
all_samples_2, _ = self.evaluate(
brs.batched_rejection_sampler(proposal, target, seed=seed))
self.assertAllEqual(all_samples, all_samples_2)
def testBatchedRejectionBetaSample(self, is_static, dtype):
seed = test_util.test_seed()
# We build a rejection sampler for two beta distributions (in a batch): a
# beta(2, 5) and a beta(2, 2). Eyeballing an image on the wikipedia page,
# these are upper bounded by rectangles of heights 2.5 and 1.6 respectively.
numpy_dtype = dtype_util.as_numpy_dtype(dtype)
alpha = np.array([2.], dtype=numpy_dtype)
beta = np.array([5., 2.], dtype=numpy_dtype)
upper_bounds = tf.constant([2.5, 1.6], dtype=dtype)
samples_per_distribution = 10000
target_fn = tfd.Beta(alpha, beta).prob
def proposal_fn(seed):
# Test static and dynamic shape of proposed samples.
uniform_samples = self.maybe_static(
tf.random.uniform([samples_per_distribution, 2],
seed=seed,
dtype=dtype), is_static)
return uniform_samples, tf.ones_like(
uniform_samples) * upper_bounds
all_samples, _ = self.evaluate(
brs.batched_rejection_sampler(proposal_fn,
target_fn,
seed=seed,
dtype=dtype))
for i in range(beta.shape[0]):
samples = all_samples[:, i]
ks, _ = sp_stats.kstest(samples, sp_stats.beta(alpha, beta[i]).cdf)
self.assertLess(ks, 0.02)
if tf.executing_eagerly():
tf.random.set_seed(seed)
# Check for reproducibility.
all_samples_2, _ = self.evaluate(
brs.batched_rejection_sampler(proposal_fn,
target_fn,
seed=seed,
dtype=dtype))
self.assertAllEqual(all_samples, all_samples_2)
def _log_concave_rejection_sampler(
mode,
prob_fn,
dtype,
sample_shape=(),
distribution_minimum=None,
distribution_maximum=None,
seed=None):
"""Utility for rejection sampling from log-concave discrete distributions.
This utility constructs an easy-to-sample-from upper bound for a discrete
univariate log-concave distribution (for discrete univariate distributions, a
necessary and sufficient condition is p_k^2 >= p_{k-1} p_{k+1} for all k).
The method requires that the mode of the distribution is known. While a better
method can likely be derived for any given distribution, this method is
general and easy to implement. The expected number of iterations is bounded by
4+m, where m is the probability of the mode. For details, see [(Devroye,
1979)][1].
Args:
mode: Tensor, the mode[s] of the [batch of] distribution[s].
prob_fn: Python callable, counts -> prob(counts).
dtype: DType of the generated samples.
sample_shape: 0D or 1D `int32` `Tensor`. Shape of the generated samples.
distribution_minimum: Tensor of type `dtype`. The minimum value
taken by the distribution. The `prob` method will only be called on values
greater than equal to the specified minimum. The shape must broadcast with
the batch shape of the distribution. If unspecified, the domain is treated
as unbounded below.
distribution_maximum: Tensor of type `dtype`. The maximum value
taken by the distribution. See `distribution_minimum` for details.
seed: Python integer or `Tensor` instance, for seeding PRNG.
Returns:
samples: a `Tensor` with prepended dimensions `sample_shape`.
#### References
[1] Luc Devroye. A Simple Generator for Discrete Log-Concave
Distributions. Computing, 1987.
[2] Dillon et al. TensorFlow Distributions. 2017.
https://arxiv.org/abs/1711.10604
"""
mode = tf.broadcast_to(
mode, tf.concat([sample_shape, prefer_static.shape(mode)], axis=0))
mode_height = prob_fn(mode)
mode_shape = prefer_static.shape(mode)
top_width = 1. + mode_height / 2. # w in ref [1].
top_fraction = top_width / (1 + top_width)
exponential_distribution = exponential.Exponential(
rate=tf.constant(1., dtype=dtype)) # E in ref [1].
if distribution_minimum is None:
distribution_minimum = tf.constant(-np.inf, dtype)
if distribution_maximum is None:
distribution_maximum = tf.constant(np.inf, dtype)
def proposal(seed):
"""Proposal for log-concave rejection sampler."""
(top_lobe_fractions_seed,
exponential_samples_seed,
top_selector_seed,
rademacher_seed) = samplers.split_seed(
seed, n=4, salt='log_concave_rejection_sampler_proposal')
top_lobe_fractions = samplers.uniform(
mode_shape, seed=top_lobe_fractions_seed, dtype=dtype) # V in ref [1].
top_offsets = top_lobe_fractions * top_width / mode_height
exponential_samples = exponential_distribution.sample(
mode_shape, seed=exponential_samples_seed) # E in ref [1].
exponential_height = (exponential_distribution.prob(exponential_samples) *
mode_height)
exponential_offsets = (top_width + exponential_samples) / mode_height
top_selector = samplers.uniform(
mode_shape, seed=top_selector_seed, dtype=dtype) # U in ref [1].
on_top_mask = tf.less_equal(top_selector, top_fraction)
unsigned_offsets = tf.where(on_top_mask, top_offsets, exponential_offsets)
offsets = tf.round(
tfp_random.rademacher(
mode_shape, seed=rademacher_seed, dtype=dtype) *
unsigned_offsets)
potential_samples = mode + offsets
envelope_height = tf.where(on_top_mask, mode_height, exponential_height)
return potential_samples, envelope_height
def target(values):
in_range_mask = (
(values >= distribution_minimum) & (values <= distribution_maximum))
in_range_values = tf.where(in_range_mask, values, 0.)
return tf.where(in_range_mask, prob_fn(in_range_values), 0.)
return tf.stop_gradient(
batched_rejection_sampler.batched_rejection_sampler(
proposal, target, seed, dtype=dtype)[0]) # Discard `num_iters`.