def reduce_dimensions(predictions, labels):
    """Reduce dimensions for high-dimensional predictions and labels."""
    # We will treat first dimensions as batch. One example are video frames.
    if len(predictions.get_shape()) > 5:
      predictions_shape = common_layers.shape_list(predictions)
      predictions = tf.reshape(
          predictions, [predictions_shape[0], predictions_shape[1], -1,
                        predictions_shape[-1]])
      labels_shape = common_layers.shape_list(labels)
      labels = tf.reshape(
          labels, [labels_shape[0], labels_shape[1], -1])
    return predictions, labels
Reduce dimensions for high-dimensional predictions and labels.
tensor2tensor/utils/metrics.py