code: def test_sample(self, data):
      batch_shape = data.draw(tfp_hps.shapes())
      bcast_arg, dist_batch_shp = data.draw(
          tfp_hps.broadcasting_shapes(batch_shape, 2)
      )
  
      underlying = normal.Normal(
          loc=tf.reshape(
              tf.range(
                  float(
                      np.prod(tensorshape_util.as_list(dist_batch_shp))
                  )
              ),
              dist_batch_shp,
          ),
          scale=0.01,
      )
  
      if not self.is_static_shape:
          bcast_arg = tf.Variable(bcast_arg)
          self.evaluate(bcast_arg.initializer)
      dist = batch_broadcast.BatchBroadcast(underlying, bcast_arg)
      sample_shape = data.draw(
          hps.one_of(hps.integers(0, 13), tfp_hps.shapes())
      )
      sample_batch_event = tf.concat(
          [
              np.int32(sample_shape).reshape([-1]),
              batch_shape,
              dist.event_shape_tensor(),
          ],
          axis=0,
      )
      sample = dist.sample(
          sample_shape,
          seed=test_util.test_seed(sampler_type="stateless"),
      )
      if self.is_static_shape:
          self.assertEqual(
              tf.TensorShape(self.evaluate(sample_batch_event)),
              sample.shape,
          )
      self.assertAllEqual(sample_batch_event, tf.shape(sample))
      # Since the `loc` of the underlying is simply 0...n-1 (reshaped), and the
      # scale is extremely small, then we can verify that these locations are
      # effectively broadcast out to the full batch shape when sampling.
      self.assertAllClose(
          tf.broadcast_to(dist.distribution.loc, sample_batch_event),
          sample,
          atol=0.1,
      )
  
      # Check that `sample_and_log_prob` also gives a correctly-shaped sample
      # with correct log-prob.
      sample2, lp = dist.experimental_sample_and_log_prob(
          sample_shape,
          seed=test_util.test_seed(sampler_type="stateless"),
      )
      if self.is_static_shape:
          self.assertEqual(
              tf.TensorShape(self.evaluate(sample_batch_event)),
              sample2.shape,
          )
      self.assertAllEqual(sample_batch_event, tf.shape(sample2))
      self.assertAllClose(lp, dist.log_prob(sample2))
  
imports: import hypothesis
  import hypothesis.strategies
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import batch_broadcast
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import hypothesis_testlib
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.distributions import von_mises_fisher
  from tensorflow_probability.python.distributions import wishart
  from tensorflow_probability.python.internal import hypothesis_testlib
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.random import random_ops
label_desc: Use image node_id 2 for calling the _BatchBroadcastTest obj's underlying member method code with example usage: obj.test_sample(data) without return types