code: def testMode(self):
      temperature = 1.0
      logits = [0.3, 0.1, 0.4]
      dist = gumbel_softmax.GumbelSoftmax(
          temperature, logits, validate_args=True
      )
      self.assertAllEqual(
          self.evaluate(dist.mode()),
          self.evaluate(tf.constant([0, 0, 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 3 for calling the GumbelSoftmaxTest obj's underlying member method code with example usage: obj.testMode() without return types