code: def testPmf(self):
      # check that probability of samples correspond to their class probabilities
      logits = self._rng.random_sample(size=(8, 2, 10))
      prob = np.exp(logits) / np.sum(
          np.exp(logits), axis=-1, keepdims=True
      )
      dist = oc.OneHotCategorical(logits=logits, validate_args=True)
      np_sample = self.evaluate(dist.sample(seed=test_util.test_seed()))
      np_prob = self.evaluate(dist.prob(np_sample))
      expected_prob = prob[np_sample.astype(np.bool_)]
      self.assertAllClose(expected_prob, np_prob.flatten())
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import onehot_categorical
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 12 for calling the OneHotCategoricalTest obj's underlying member method code with example usage: obj.testPmf() without return types