code: def sample_with_temperature(x, dim, temperature=1.0, dtype=tf.int32, name=None):
    
    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)
docs: 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.