code: def test_sample_distribution(self):
      probs = np.random.rand(4) + 0.001
      probs = probs / np.sum(probs)
  
      def model():
          i = yield Root(
              sample_dist_lib.Sample(
                  categorical.Categorical(probs=probs, dtype=tf.int32),
                  sample_shape=[2],
              )
          )
          # Note use of scalar `sample_shape` to test expansion of shape
          # to vector.
          j = yield sample_dist_lib.Sample(
              categorical.Categorical(probs=probs, dtype=tf.int32),
              sample_shape=2,
          )
  
      dist = marginalize.MarginalizableJointDistributionCoroutine(model)
  
      ptt = tf.exp(dist.marginalized_log_prob(["tabulate", "tabulate"]))
      ptm = tf.exp(
          dist.marginalized_log_prob(["tabulate", "marginalize"])
      )
      pmt = tf.exp(
          dist.marginalized_log_prob(["marginalize", "tabulate"])
      )
      pmm = tf.exp(
          dist.marginalized_log_prob(["marginalize", "marginalize"])
      )
      self.assertEqual(ptt.shape, [4, 4, 4, 4])
      self.assertEqual(ptm.shape, [4, 4])
      self.assertEqual(pmt.shape, [4, 4])
      self.assertEqual(pmm.shape, [])
  
imports: import functools
  import itertools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import hidden_markov_model
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  import tensorflow_probability.python.experimental.marginalize
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 4 for calling the _MarginalizeTest obj's underlying member method code with example usage: obj.test_sample_distribution() without return types