code: def logsoftmax_sample(
      log_probs, temperature=1.0
  ):  # pylint: disable=invalid-name
      """Returns a sample from a log-softmax output, with temperature.
  
      Args:
        log_probs: Logarithms of probabilities (often coming from LogSoftmax)
        temperature: For scaling before sampling (1.0 = default, 0.0 = pick argmax)
      """
      # This is equivalent to sampling from a softmax with temperature.
      u = np.random.uniform(
          low=1e-6, high=1.0 - 1e-6, size=log_probs.shape
      )
      g = -np.log(-np.log(u))
      return np.argmax(log_probs + g * temperature, axis=-1)
  
imports: from absl import logging
  import numpy
  import tensorflow
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 20 for calling a global function with example usage: logsoftmax_sample(log_probs, temperature) and returns: np