//  [python]
def _slow_greedy_infer_tpu(self, features, decode_length):
        """A slow greedy inference method on TPU.

    Quadratic time in decode_length.

    Args:
      features: An map of string to `Tensor`.
      decode_length: An integer, how many additional timesteps to decode.

    Returns:
      A dict of decoding results {
          "outputs": integer `Tensor` of decoded ids of shape
              [batch_size, <= decode_length] if beam_size == 1 or
              [batch_size, top_beams, <= decode_length]
          "scores": None
          "logits": `Tensor` of shape [batch_size, time, 1, 1, vocab_size].
          "losses": a dictionary: {loss-name (string): floating point `Scalar`}
      }
    """
        if not features:
            features = {}
        inputs_old = None
        if 'inputs' in features and len(features['inputs'].shape) < 4:
            inputs_old = features['inputs']
            features['inputs'] = tf.expand_dims(features['inputs'], 2)
        if not self.has_input:
            # Prepare partial targets.
            # In either features["inputs"] or features["targets"].
            # We force the outputs to begin with these sequences.
            partial_targets = features.get('inputs')
            if partial_targets is None:
                partial_targets = features['targets']
            features['partial_targets'] = tf.cast(partial_targets, tf.int64)
        # Save the targets in a var and reassign it after the tf.while loop to avoid
        # having targets being in a 'while' frame. This ensures targets when used
        # in metric functions stays in the same frame as other vars.
        targets_old = features.get('targets', None)

        target_modality = self._problem_hparams.modality['targets']

        def infer_step(i, recent_output, recent_logits, unused_loss):
            """Inference step."""
            if not tf.executing_eagerly():
                recent_output.set_shape([None, None, None, 1])
            padded = tf.pad(recent_output, [[0, 0], [0, 1], [0, 0], [0, 0]])
            features['targets'] = padded
            # This is inefficient in that it generates samples at all timesteps,
            # not just the last one, except if target_modality is pointwise.
            features['decode_loop_step'] = i
            samples, logits, losses = self.sample(features)
            # Concatenate the already-generated recent_output with last timestep
            # of the newly-generated samples.z
            top = self._hparams.top.get(
                'targets', modalities.get_top(target_modality)
            )
            if getattr(top, 'pointwise', False):
                cur_sample = samples[:, -1, :, :]
            else:
                cur_sample = samples[:, i, :, :]
            samples = tf.transpose(recent_output, perm = [1, 0, 2, 3])
            samples = inplace_ops.alias_inplace_update(
                samples, i, tf.cast(cur_sample, tf.int64)
            )
            samples = tf.transpose(samples, perm = [1, 0, 2, 3])
            if not tf.executing_eagerly():
                samples.set_shape([None, None, None, 1])

            # Assuming we have one shard for logits.
            recent_logits = tf.transpose(recent_logits, perm = [1, 0, 2, 3, 4])
            recent_logits = inplace_ops.alias_inplace_update(
                recent_logits, i, tf.squeeze(logits[:, -1:], axis = 1)
            )
            logits = tf.transpose(recent_logits, perm = [1, 0, 2, 3, 4])
            loss = sum([l for l in losses.values() if l is not None])
            return i + 1, samples, logits, loss

        # Create an initial output tensor. This will be passed
        # to the infer_step, which adds one timestep at every iteration.
        if 'partial_targets' in features:
            initial_output = tf.cast(features['partial_targets'], tf.int64)
            while len(initial_output.get_shape().as_list()) < 4:
                initial_output = tf.expand_dims(initial_output, 2)
            batch_size = common_layers.shape_list(initial_output)[0]
        else:
            batch_size = common_layers.shape_list(features['inputs'])[0]
            initial_output = tf.zeros((batch_size, 0, 1, 1), dtype = tf.int64)
        # Hack: foldl complains when the output shape is less specified than the
        # input shape, so we confuse it about the input shape.
        initial_output = tf.slice(
            initial_output,
            [0, 0, 0, 0],
            common_layers.shape_list(initial_output),
        )
        target_modality = self._problem_hparams.modality['targets']
        if (
            target_modality == modalities.ModalityType.CLASS_LABEL
            or self._problem_hparams.get('regression_targets')
        ):
            decode_length = 1
        else:
            if 'partial_targets' in features:
                prefix_length = common_layers.shape_list(
                    features['partial_targets']
                )[1]
            else:
                prefix_length = common_layers.shape_list(features['inputs'])[1]
            decode_length = prefix_length + decode_length

        # Initial values of result, logits and loss.
        result = tf.concat(
            [
                initial_output,
                tf.zeros([batch_size, decode_length, 1, 1], tf.int64),
            ],
            axis = 1,
        )
        # tensor padded to [batch_size, decode_length, 1, 1, vocab_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
        logits = tf.zeros((batch_size, decode_length, 1, 1, vocab_size))
        if not tf.executing_eagerly():
            logits.set_shape([None, None, None, None, None])
        loss = 0.0

        def while_exit_cond(
            i, result, logits, loss
        ):  # pylint: disable=unused-argument
            """Exit the loop either if reach decode_length or EOS."""
            not_overflow = i < decode_length

            if self._problem_hparams.stop_at_eos:

                def fn_not_eos():
                    # Check if the last predicted element is a EOS
                    return tf.reduce_any(
                        tf.not_equal(
                            tf.squeeze(result[:, -1, :, :]), text_encoder.EOS_ID
                        )
                    )

                not_eos = tf.cond(
                    # We only check for early stopping if there is at least 1 element (
                    # otherwise not_eos will crash).
                    tf.not_equal(i, 0),
                    fn_not_eos,
                    lambda: True,
                )

                return tf.cond(
                    tf.equal(batch_size, 1),
                    # If batch_size == 1, we check EOS for early stopping.
                    lambda: tf.logical_and(not_overflow, not_eos),
                    # Else, just wait for max length
                    lambda: not_overflow,
                )
            return not_overflow

        _, result, logits, loss = tf.while_loop(
            while_exit_cond,
            infer_step,
            [tf.constant(0), result, logits, loss],
            shape_invariants = [
                tf.TensorShape([]),
                tf.TensorShape([batch_size, decode_length, 1, 1]),
                tf.TensorShape([batch_size, decode_length, 1, 1, vocab_size]),
                tf.TensorShape([]),
            ],
            back_prop = False,
            parallel_iterations = 1,
        )
        if inputs_old is not None:  # Restore to not confuse Estimator.
            features['inputs'] = inputs_old
        # Reassign targets back to the previous value.
        if targets_old is not None:
            features['targets'] = targets_old
        losses = {'training': loss}
        if 'partial_targets' in features:
            partial_target_length = common_layers.shape_list(
                features['partial_targets']
            )[1]
            result = tf.slice(
                result, [0, partial_target_length, 0, 0], [-1, -1, -1, -1]
            )
        return {
            'outputs': result,
            'scores': None,
            'logits': logits,
            'losses': losses,
        }