code: def test_unit_sample_shape(self):
      @jdab.JointDistributionCoroutineAutoBatched
      def dist():
          x = yield normal.Normal(
              loc=tf.zeros([3]), scale=1.0, name="x"
          )
          yield bernoulli.Bernoulli(
              logits=tf.einsum("n->", x), name="y"
          )
  
      for sample_shape in [(), 1, [1], [1, 1], [2]]:
          self.assertAllEqual(
              dist.log_prob(
                  dist.sample(sample_shape, seed=test_util.test_seed())
              ).shape,
              np.reshape(sample_shape, [-1]),
          )
  
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 15 for calling the JointDistributionAutoBatchedTest obj's underlying member method code with example usage: obj.test_unit_sample_shape() without return types