code: def sample(self, features):
      
      logits, losses = self(features)  # pylint: disable=not-callable
      if self._target_modality_is_real:
        return logits, logits, losses  # Raw numbers returned from real modality.
      if self.hparams.sampling_method == "argmax":
        samples = tf.argmax(logits, axis=-1)
      else:
        assert self.hparams.sampling_method == "random"
  
        def multinomial_squeeze(logits, temperature=1.0):
          logits_shape = common_layers.shape_list(logits)
          reshaped_logits = (
              tf.reshape(logits, [-1, logits_shape[-1]]) / temperature)
          choices = tf.multinomial(reshaped_logits, 1)
          choices = tf.reshape(choices, logits_shape[:-1])
          return choices
  
        samples = multinomial_squeeze(logits, self.hparams.sampling_temp)
  
      return samples, logits, losses
docs: Run the model and extract samples.
  
      Args:
        features: an map of string to `Tensor`.
  
      Returns:
         samples: an integer `Tensor`.
         logits: a list of `Tensor`s, one per datashard.
         losses: a dictionary: {loss-name (string): floating point `Scalar`}.