code: def gumbel_softmax(
      logits,
      temperature: float = 1.0,
      eps: float = 1e-10,
      hard=True,
      use_np_gumbel: bool = True,
  ):
      r"""Perform the gumbel-softmax trick to generate differentiable one-hot vectors from the input
      logits.
  
      Here, the gumbel distribution is
  
      Gumbel(\alpha) = -log (-log U) + \log \alpha, in which U is the uniform(0, 1) distribution.
  
      A nice property of Gumbel is:
  
      \argmax({Gumbel(\alpha_i)}) \sim multinomial(\alpha_i)
  
      The Gumbel-Softmax trick is to use the softmax + straight-through estimator to produce
      one-hot vectors that represent the sampling result.
  
      References:
  
          1. https://en.wikipedia.org/wiki/Gumbel_distribution
          2. [ICLR2017] Categorical Reparameterization with Gumbel-Softmax
  
      Parameters
      ----------
      logits
          Logits. Shape (..., V)
      temperature
          The temperature that controls the
      eps
          The eps for stability of gradient
      hard
          Whether to use the straight-through estimator to produce one-hot vectors.
      use_np_gumbel
          Whether to use the random.gumble operator
  
      Returns
      -------
      ret
          The returned output. Shape (..., V)
      """
      # TODO(sxjscience) Investigate the impact of random.gumbel:
      #  Actually, random.gumble has no eps and may have problem in calculating the gradient.
      if use_np_gumbel:
          gumbels = np.random.gumbel(np.zeros_like(logits))
      else:
          u = np.random.uniform(np.zeros_like(logits), 1)
          gumbels = -np.log(-np.log(u + eps) + eps)
      y = npx.softmax((gumbels + logits) / temperature, axis=-1)
      if hard:
          y_hard = np.max(y, axis=-1, keepdims=True) == y
          y_hard = npx.stop_gradient(y_hard - y) + y
          return y_hard
      else:
          return y
  
file_path: gluon-nlp/src/gluonnlp/op.py
imports: import mxnet
  from mxnet import np, npx
  import math
  from mxnet import use_np
label_desc: Use image node_id 4 for calling a global function with example usage: gumbel_softmax(logits, temperature, eps, hard, use_np_gumbel) and returns: y_hard, y