def _prepare_inputs_for_decode(self, features):
    """Prepare inputs for decoding.

    Args:
      features: A map of string to model features.

    Returns:
      Inputs after fixing shape and applying modality.
    """
    dp = self._data_parallelism
    hparams = self._hparams
    inputs = features["inputs"]
    # TODO(llion): Clean up this reshaping logic.
    inputs = tf.expand_dims(inputs, axis=1)
    if len(inputs.shape) < 5:
      inputs = tf.expand_dims(inputs, axis=4)
    s = common_layers.shape_list(inputs)
    inputs = tf.reshape(inputs, [s[0] * s[1], s[2], s[3], s[4]])
    # _shard_features called to ensure that the variable names match
    inputs = self._shard_features({"inputs": inputs})["inputs"]
    input_modality = self._problem_hparams.modality["inputs"]
    input_vocab_size = self._problem_hparams.vocab_size["inputs"]
    if input_vocab_size is not None and hasattr(hparams, "vocab_divisor"):
      input_vocab_size += (-input_vocab_size) % hparams.vocab_divisor
    modality_name = hparams.name.get("inputs",
                                     modalities.get_name(input_modality))(
                                         hparams, input_vocab_size)
    with tf.variable_scope(modality_name):
      bottom = hparams.bottom.get("inputs",
                                  modalities.get_bottom(input_modality))
      inputs = dp(bottom, inputs, hparams, input_vocab_size)
    return inputs
Prepare inputs for decoding.

    Args:
      features: A map of string to model features.

    Returns:
      Inputs after fixing shape and applying modality.
    
tensor2tensor/models/transformer.py