code: def testLogPMFBroadcasting(self):
      # 1 x 2 x 2
      histograms = np.array([[[0.2, 0.8], [0.4, 0.6]]])
      dist = categorical.Categorical(
          tf.math.log(histograms) - 50.0, validate_args=True
      )
  
      prob = dist.prob(1)
      self.assertAllClose([[0.8, 0.6]], self.evaluate(prob))
  
      prob = dist.prob([1])
      self.assertAllClose([[0.8, 0.6]], self.evaluate(prob))
  
      prob = dist.prob([0, 1])
      self.assertAllClose([[0.2, 0.6]], self.evaluate(prob))
  
      prob = dist.prob([[0, 1]])
      self.assertAllClose([[0.2, 0.6]], self.evaluate(prob))
  
      prob = dist.prob([[[0, 1]]])
      self.assertAllClose([[[0.2, 0.6]]], self.evaluate(prob))
  
      prob = dist.prob([[1, 0], [0, 1]])
      self.assertAllClose([[0.8, 0.4], [0.2, 0.6]], self.evaluate(prob))
  
      prob = dist.prob([[[1, 1], [1, 0]], [[1, 0], [0, 1]]])
      self.assertAllClose(
          [[[0.8, 0.6], [0.8, 0.4]], [[0.8, 0.4], [0.2, 0.6]]],
          self.evaluate(prob),
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  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 26 for calling the CategoricalTest obj's underlying member method code with example usage: obj.testLogPMFBroadcasting() without return types