class_docstr: Tests ActionDistribution classes.
docstr: Tests the GumbelSoftmax ActionDistribution (tf + eager only).
docstr_tok: ["Tests","the","GumbelSoftmax","ActionDistribution","(","tf","+","eager","only",")","."]
code: def test_gumbel_softmax(self):
      """Tests the GumbelSoftmax ActionDistribution (tf + eager only)."""
      for fw, sess in framework_iterator(
          frameworks=("tf2", "tf"), session=True
      ):
          batch_size = 1000
          num_categories = 5
          input_space = Box(
              -1.0, 1.0, shape=(batch_size, num_categories)
          )
          input_space.seed(42)
  
          # Batch of size=n and deterministic.
          inputs = input_space.sample()
          gumbel_softmax = GumbelSoftmax(inputs, {}, temperature=1.0)
  
          expected = softmax(inputs)
          # Sample n times, expect always mean value (deterministic draw).
          out = gumbel_softmax.deterministic_sample()
          check(out, expected)
  
          # Batch of size=n and non-deterministic -> expect roughly that
          # the max-likelihood (argmax) ints are output (most of the time).
          inputs = input_space.sample()
          gumbel_softmax = GumbelSoftmax(inputs, {}, temperature=1.0)
          expected_mean = np.mean(np.argmax(inputs, -1)).astype(
              np.float32
          )
          outs = gumbel_softmax.sample()
          if sess:
              outs = sess.run(outs)
          check(np.mean(np.argmax(outs, -1)), expected_mean, rtol=0.08)
  
imports: from functools import partial
  from gymnasium.spaces import Box, Dict, Tuple
  import numpy
  from scipy.stats import beta, norm
  import tree
  import unittest
  from ray.rllib.models.tf.tf_action_dist import Beta, Categorical, DiagGaussian, GumbelSoftmax, MultiActionDistribution, MultiCategorical, SquashedGaussian
  from ray.rllib.models.torch.torch_action_dist import TorchBeta, TorchCategorical, TorchDiagGaussian, TorchMultiActionDistribution, TorchMultiCategorical, TorchSquashedGaussian
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.numpy import MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT, softmax, SMALL_NUMBER, LARGE_INTEGER
  from ray.rllib.utils.test_utils import check, framework_iterator
label_desc: Use image node_id 8 for calling the TestActionDistributions obj's underlying member method code with example usage: obj.test_gumbel_softmax() without return types