code: def testPmfCountsStretchedInBroadcastWhenSameRank(self):
      p = [[0.1, 0.9], [0.7, 0.3]]
      counts = [[1.0, 0]]
      pmf = multinomial.Multinomial(
          total_count=1.0, probs=p, validate_args=True
      ).prob(counts)
      self.assertAllClose(self.evaluate(pmf), [0.1, 0.7])
      self.assertAllEqual([2], pmf.shape)
  
imports: import unittest
  from absl.testing import parameterized
  import hypothesis
  from hypothesis import strategies
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import binomial
  from tensorflow_probability.python.distributions import hypothesis_testlib
  from tensorflow_probability.python.distributions import multinomial
  from tensorflow_probability.python.distributions.internal import statistical_testing
  from tensorflow_probability.python.internal import hypothesis_testlib
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 14 for calling the MultinomialTest obj's underlying member method code with example usage: obj.testPmfCountsStretchedInBroadcastWhenSameRank() without return types