code: def test_unit_sample_shape_avoids_vectorization(self):
      xs = (
          []
      )  # Collect (possibly symbolic) Tensors sampled inside the model.
  
      @jdab.JointDistributionCoroutineAutoBatched
      def dist():
          x = yield normal.Normal(0.0, 1.0, name="x")
          xs.append(x)
  
      # Try sampling with a variety of unit sample shapes.
      self.assertEqual(
          [1],
          dist.sample(
              1, seed=test_util.test_seed(sampler_type="seedless")
          ).x.shape,
      )
      self.assertEqual(
          [1],
          dist.sample(
              [1], seed=test_util.test_seed(sampler_type="seedless")
          ).x.shape,
      )
      self.assertEqual(
          [1, 1],
          dist.sample(
              [1, 1], seed=test_util.test_seed(sampler_type="seedless")
          ).x.shape,
      )
  
      # Check that the model only ever saw the trivial sample shape.
      for x in xs:
          self.assertEqual(x.shape, [])
  
imports: import collections
  import os
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector_test_util
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import multinomial
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 14 for calling the JointDistributionAutoBatchedTest obj's underlying member method code with example usage: obj.test_unit_sample_shape_avoids_vectorization() without return types