code: def testPmfShapeCountsStretchedN(self):
      with self.cached_session():
          # [2, 2, 2]
          p = [[[0.1, 0.9], [0.1, 0.9]], [[0.7, 0.3], [0.7, 0.3]]]
          # [2, 2]
          n = [[3.0, 3], [3, 3]]
          # [2]
          counts = [2.0, 1]
          pmf = multinomial.Multinomial(total_count=n, probs=p).prob(
              counts
          )
          self.evaluate(pmf)
          self.assertEqual(pmf.get_shape(), (2, 2))
  
imports: import numpy
  from tensorflow.python.eager import backprop
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops.distributions import multinomial
  from tensorflow.python.platform import test
label_desc: Use image node_id 16 for calling the MultinomialTest obj's underlying member method code with example usage: obj.testPmfShapeCountsStretchedN() without return types