def top(self, body_output, features):
    """Adds additional dimensions and then calls super class implementation."""
    exp_features = features
    for feat in body_output.keys():
      while len(body_output[feat].shape) < 4:
        logging.warning("Expanding body output %s...", feat)
        body_output[feat] = tf.expand_dims(body_output[feat], -2)
      if feat in exp_features:
        while len(exp_features[feat].shape) < 4:
          exp_features[feat] = tf.expand_dims(exp_features[feat], -1)
          logging.warning("Expanding feature %s...", feat)
    return super(TransformerSeq2edits, self).top(body_output, exp_features)
Adds additional dimensions and then calls super class implementation.
tensor2tensor/models/research/transformer_seq2edits.py