code: def testLogPMFBroadcasting(self):
      with self.cached_session():
          # 1 x 2 x 2
          histograms = [[[0.2, 0.8], [0.4, 0.6]]]
          dist = categorical.Categorical(
              math_ops.log(histograms) - 50.0
          )
  
          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
  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_util
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gradients_impl
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops.distributions import categorical
  from tensorflow.python.ops.distributions import kullback_leibler
  from tensorflow.python.ops.distributions import normal
  from tensorflow.python.platform import test
label_desc: Use image node_id 21 for calling the CategoricalTest obj's underlying member method code with example usage: obj.testLogPMFBroadcasting() without return types