def image_summary(self, name, image_logits, max_outputs=1):
    """Helper for image summaries that are safe on TPU."""
    if len(image_logits.get_shape()) != 5:
      tf.logging.info("Not generating image summary, maybe not an image.")
      return
    return tf.summary.image(
        name,
        common_layers.tpu_safe_image_summary(tf.argmax(image_logits, -1)),
        max_outputs=max_outputs)
Helper for image summaries that are safe on TPU.
tensor2tensor/models/research/autoencoders.py