class_docstr: Image-to-image transformer predicting blocks of the output in parallel.
class_docstr_tok: ["Image-to-image","transformer","predicting","blocks","of","the","output","in","parallel","."]
code: def loss(self, logits, features):
      assert self._hparams.block_size > 0
  
      if self._hparams.mode == tf_estimator.ModeKeys.PREDICT:
          return 0.0
  
      def shift_left_2d(x, k):
          return tf.pad(x, [[0, 0], [0, k]])[:, k:]
  
      def shift_left_4d_raster_scan(x, k):
          batch_size = common_layers.shape_list(x)[0]
          return tf.reshape(
              shift_left_2d(tf.reshape(x, [batch_size, -1]), k),
              tf.shape(x),
          )
  
      targets = features["targets"]
      assert len(targets.shape) == 4
  
      targets = tf.stack(
          [
              shift_left_4d_raster_scan(targets, i)
              for i in range(self._hparams.block_size)
          ],
          axis=4,
      )
  
      if (
          self._hparams.mode == tf_estimator.ModeKeys.TRAIN
          or self._hparams.mode == tf_estimator.ModeKeys.EVAL
      ):
          assert "block_index" in features
          targets = targets[:, :, :, :, features["block_index"]]
  
      features["targets"] = targets
  
      loss = super(Img2imgTransformerBlockParallel, self).loss(
          logits, features
      )
  
      if self._hparams.mode == tf_estimator.ModeKeys.TRAIN:
          k = features["block_index"]
          loss_num, loss_den = loss
          loss_val = loss_num / loss_den
          for i in range(self._hparams.block_size):
              # Hack: if you report a loss of NaN, TensorBoard will plot a point at
              # the previous value without a connecting line. This is used here to
              # separate out the training losses by block index.
              one_or_nan = tf.cond(
                  tf.equal(k, i), lambda: 1.0, lambda: float("nan")
              )
              tf.summary.scalar(
                  "block_index_%d" % i,
                  one_or_nan * loss_val,
                  family="losses",
              )
  
      return loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import numpy
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import modalities
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 3 for calling the Img2imgTransformerBlockParallel obj's underlying member method code with example usage: obj.loss(logits, features) and returns: loss, int, tf, tf