code: def _test_vectorization(self, dist_name, dist):
      seed = test_util.test_seed()
  
      # TODO(b/171752261): New stateless samplers don't work with pfor.
      enable_auto_vectorized_sampling = False
  
      num_samples = 3
      if (
          not enable_auto_vectorized_sampling
          or dist_name in SAMPLE_AUTOVECTORIZATION_IS_BROKEN
      ):
          sample = self.evaluate(dist.sample(num_samples, seed=seed))
      else:
          sample = self.evaluate(
              tf.vectorized_map(
                  lambda i: dist.sample(seed=seed),
                  tf.range(num_samples),
                  fallback_to_while_loop=False,
              )
          )
      hp.note("Drew samples {}".format(sample))
  
      if dist_name not in LOGPROB_AUTOVECTORIZATION_IS_BROKEN:
          pfor_lp = tf.vectorized_map(
              dist.log_prob,
              tf.convert_to_tensor(sample),
              fallback_to_while_loop=False,
          )
          batch_lp = dist.log_prob(sample)
          pfor_lp_, batch_lp_ = self.evaluate((pfor_lp, batch_lp))
          self.assertAllClose(
              pfor_lp_,
              batch_lp_,
              atol=VECTORIZED_LOGPROB_ATOL[dist_name],
              rtol=VECTORIZED_LOGPROB_RTOL[dist_name],
          )
  
imports: import collections
  from absl.testing import parameterized
  import hypothesis
  from hypothesis import strategies
  import numpy
  import six
  import tensorflow.compat.v2
  from tensorflow_probability.python import distributions
  from tensorflow_probability.python.distributions import hypothesis_testlib
  from tensorflow_probability.python.internal import hypothesis_testlib
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the DistributionsWorkWithAutoVectorizationTest obj's underlying member method code with example usage: obj._test_vectorization(dist_name, dist) without return types