original_docstring: Sample from the Gumbel-Softmax distribution and optionally discretize.
      Args:
        logits: [batch_size, n_class] unnormalized log-probs
        temperature: non-negative scalar
        hard: if True, take argmax, but differentiate w.r.t. soft sample y
      Returns:
        [batch_size, n_class] sample from the Gumbel-Softmax distribution.
        If hard=True, then the returned sample will be one-hot, otherwise it will
        be a probabilitiy distribution that sums to 1 across classes
      
docstring: Sample from the Gumbel-Softmax distribution and optionally discretize.
docstring_tokens: Sample from the Gumbel - Softmax distribution and optionally discretize .
code: def gumbel_softmax(logits, temperature=0.5, hard=False):
      """Sample from the Gumbel-Softmax distribution and optionally discretize.
      Args:
        logits: [batch_size, n_class] unnormalized log-probs
        temperature: non-negative scalar
        hard: if True, take argmax, but differentiate w.r.t. soft sample y
      Returns:
        [batch_size, n_class] sample from the Gumbel-Softmax distribution.
        If hard=True, then the returned sample will be one-hot, otherwise it will
        be a probabilitiy distribution that sums to 1 across classes
      """
      y = gumbel_softmax_sample(logits, temperature)
      if hard:
          y_hard = onehot_from_logits(y)
          y = (y_hard - y).detach() + y
      return y
code_tokens: def gumbel_softmax ( logits , temperature = 0.5 , hard = False ) : y = gumbel_softmax_sample ( logits , temperature ) if hard : y_hard = onehot_from_logits ( y ) y = ( y_hard - y ) . detach ( ) + y return y
short_docstring: Sample from the Gumbel-Softmax distribution and optionally discretize.
short_docstring_tokens: Sample from the Gumbel - Softmax distribution and optionally discretize .
comment: """Sample from the Gumbel-Softmax distribution and optionally discretize.
      Args:
        logits: [batch_size, n_class] unnormalized log-probs
        temperature: non-negative scalar
        hard: if True, take argmax, but differentiate w.r.t. soft sample y
      Returns:
        [batch_size, n_class] sample from the Gumbel-Softmax distribution.
        If hard=True, then the returned sample will be one-hot, otherwise it will
        be a probabilitiy distribution that sums to 1 across classes
      """