Slow version of Beam search decoding.

    Quadratic time in decode_length.

    Args:
      features: an map of string to `Tensor`
      decode_length: an integer.  How many additional timesteps to decode.
      beam_size: number of beams.
      top_beams: an integer. How many of the beams to return.
      alpha: Float that controls the length penalty. larger the alpha, stronger
        the preference for longer translations.
      use_tpu: A bool, whether to do slow beam decode on TPU.

    Returns:
      samples: an integer `Tensor`. Top samples from the beam search.

    Raises:
      NotImplementedError: If use_tpu is set to true.

def _beam_decode_slow(self, features, decode_length, beam_size, top_beams,
                        alpha, use_tpu=False):
    """Slow version of Beam search decoding.

    Quadratic time in decode_length.

    Args:
      features: an map of string to `Tensor`
      decode_length: an integer.  How many additional timesteps to decode.
      beam_size: number of beams.
      top_beams: an integer. How many of the beams to return.
      alpha: Float that controls the length penalty. larger the alpha, stronger
        the preference for longer translations.
      use_tpu: A bool, whether to do slow beam decode on TPU.

    Returns:
      samples: an integer `Tensor`. Top samples from the beam search.

    Raises:
      NotImplementedError: If use_tpu is set to true.
    """
    batch_size = common_layers.shape_list(features["inputs"])[0]

    def symbols_to_logits_fn(ids, i=None):
      """Go from ids to logits."""
      ids = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
      ids = tf.pad(ids[:, 1:], [[0, 0], [0, 1], [0, 0], [0, 0]])
      if "partial_targets" in features:
        pt = features["partial_targets"]
        pt_length = common_layers.shape_list(pt)[1]
        pt = tf.tile(pt, [1, beam_size])
        pt = tf.reshape(pt, [batch_size * beam_size, pt_length, 1, 1])
        ids = tf.concat([pt, ids], axis=1)

      features["targets"] = ids
      if i is not None:
        features["decode_loop_step"] = i
      self._coverage = None
      logits, _ = self(features)  # pylint: disable=not-callable
      # now self._coverage is a coverage tensor for the first datashard.
      # it has shape [batch_size] and contains floats between 0 and
      # source_length.
      if self._problem_hparams:
        modality = self._problem_hparams.modality["targets"]
        top = self._hparams.top.get("targets", modalities.get_top(modality))
        if getattr(top, "pointwise", False):
          return tf.squeeze(logits, axis=[1, 2, 3])
      # -1 due to the pad above.
      current_output_position = common_layers.shape_list(ids)[1] - 1
      logits = logits[:, current_output_position, :, :]
      return tf.squeeze(logits, axis=[1, 2])

    def _clone_examples_for_beam(old_feature, n):
      """Clone each example n times."""
      old_shape = common_layers.shape_list(old_feature)
      assert len(old_shape) >= 1

      # Expand the inputs in to the beam size.
      feature = tf.expand_dims(old_feature, 1)
      feature = tf.tile(feature, [1, n] + [1] * (len(old_shape) - 1))
      new_shape = common_layers.shape_list(feature)
      feature = tf.reshape(feature,
                           [new_shape[0] * new_shape[1]] + new_shape[2:])
      return feature

    initial_ids = tf.zeros([batch_size], dtype=tf.int32)

    # Clone select features multiple times to account for beam size.
    old_features = {}
    for feature_name in ["inputs", "knowledge"]:
      if feature_name not in features:
        continue
      old_features[feature_name] = features[feature_name]
      features[feature_name] = _clone_examples_for_beam(
          features[feature_name], beam_size)

    vocab_size = self._problem_hparams.vocab_size["targets"]
    if vocab_size is not None and hasattr(self._hparams, "vocab_divisor"):
      vocab_size += (-vocab_size) % self._hparams.vocab_divisor

    # Setting decode length to input length + decode_length
    if "partial_targets" not in features:
      inputs = features["inputs"]
      decode_length = (common_layers.shape_list(inputs)[1] +
                       features.get("decode_length", decode_length))
    ids, scores, _ = beam_search.beam_search(
        symbols_to_logits_fn,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        alpha,
        stop_early=(top_beams == 1),
        use_tpu=use_tpu)

    # Set features back to the unexpanded form to not to confuse the
    # Estimator!
    features.update(old_features)

    # Return `top_beams` decodings (also remove initial id from the beam search)
    # TODO(lukaszkaiser): make it work multi-problem.
    if top_beams == 1:
      samples = ids[:, 0, 1:]
    else:
      samples = ids[:, :top_beams, 1:]

    return {"outputs": samples, "scores": scores}