code: def testSample(self):
      temperature = 0.8
      logits = [0.3, 0.1, 0.4]
      dist = gumbel_softmax.GumbelSoftmax(
          temperature, logits, dtype=tf.int64, validate_args=True
      )
      actions = dist.convert_to_one_hot(dist.sample())
      self.assertEqual(actions.dtype, tf.int64)
      self.assertEqual(
          self.evaluate(tf.reduce_sum(actions, axis=-1)), 1
      )
  
file_path: agents/tf_agents/distributions/gumbel_softmax_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tf_agents.distributions import gumbel_softmax
label_desc: Use image node_id 2 for calling the GumbelSoftmaxTest obj's underlying member method code with example usage: obj.testSample() without return types