code: def testSample(self):
      probs = [[[0.2, 0.8], [0.4, 0.6]]]
      dist = oc.OneHotCategorical(
          np.log(probs) - 50.0, validate_args=True
      )
      n = 100
      samples = dist.sample(n, seed=test_util.test_seed())
      self.assertEqual(samples.dtype, tf.int32)
      sample_values = self.evaluate(samples)
      self.assertAllEqual([n, 1, 2, 2], sample_values.shape)
      self.assertFalse(np.any(sample_values < 0))
      self.assertFalse(np.any(sample_values > 1))
  
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 14 for calling the OneHotCategoricalTest obj's underlying member method code with example usage: obj.testSample() without return types