'Run the model and extract samples.
Args:
features: an map of string to `Tensor`.
last_position_only: a boolean, speed-up by computing last position only.
Returns:
samples: an integer `Tensor`.
logits: a list of `Tensor`s, one per datashard.
losses: a dictionary: {loss-name (string): floating point `Scalar`}.'

def sample(self, features, last_position_only=False):

    (sharded_logits, losses) = self.model_fn(features, False, last_position_only=last_position_only)
    if (self._hparams.sampling_method == 'argmax'):
        sharded_samples = self._data_parallelism(tf.argmax, sharded_logits, 4)
    else:
        assert (self._hparams.sampling_method == 'random')
        def _multinomial_squeeze(logits):
            reshaped_logits = tf.reshape(logits, [(-1), tf.shape(logits)[(-1)]])
            choices = tf.multinomial(reshaped_logits, 1)
            choices = tf.reshape(choices, tf.shape(logits)[:(logits.get_shape().ndims - 1)])
            return choices
        sharded_samples = self._data_parallelism(_multinomial_squeeze, sharded_logits)
    return (tf.concat(sharded_samples, 0), sharded_logits, losses)
