def sample_with_temperature(x, dim, temperature=1.0, dtype=tf.int32, name=None):
  """Either argmax or random sampling.

  Args:
    x: a Tensor.
    dim: a Dimension in x.shape.dims
    temperature: a float  0.0=argmax 1.0=random
    dtype: a tf.dtype (for the output)
    name: an optional string

  Returns:
    a Tensor with type dtype.
  """
  dim = convert_to_dimension(dim)
  with tf.name_scope(name, default_name="sample_with_temperature"):
    if temperature != 0.0:
      # gumbel trick.
      # Note: we don't want to generate 0 or 1 because:
      # * -log(-log(0)) is -infinity
      # * -log(-log(1)) is +infinity.
      # np.finfo(x.dtype.as_numpy_dtype).tiny doesn't work on bfloat16
      tiny_val = 1e-9
      g = -log(-log(
          random_uniform(
              x.mesh,
              x.shape,
              minval=tiny_val,
              maxval=1.,
              dtype=x.dtype)))
      x += g * temperature
    return argmax(x, dim, dtype, name)