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`.