code: def testBatchedRejectionBetaSample(self, is_static, dtype):
      # 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.0], dtype=numpy_dtype)
      beta = np.array([5.0, 2.0], dtype=numpy_dtype)
      upper_bounds = tf.constant([2.5, 1.6], dtype=dtype)
      samples_per_distribution = 10000
  
      target_fn = beta_lib.Beta(alpha, beta).prob
  
      def proposal_fn(seed):
          # Test static and dynamic shape of proposed samples.
          uniform_samples = self.maybe_static(
              samplers.uniform(
                  [samples_per_distribution, 2], seed=seed, dtype=dtype
              ),
              is_static,
          )
          return (
              uniform_samples,
              tf.ones_like(uniform_samples) * upper_bounds,
          )
  
      seed = test_util.test_seed(sampler_type="stateless")
  
      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)
  
      # 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)
  
file_path: probability/tensorflow_probability/python/internal/batched_rejection_sampler_test.py
imports: from absl.testing import parameterized
  import numpy
  from scipy import stats
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.internal import batched_rejection_sampler
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 4 for calling the BatchedRejectionSamplerTest obj's underlying member method code with example usage: obj.testBatchedRejectionBetaSample(is_static, dtype) and returns: uniform_samples,  unknown