class_docstr: Transformer that predicts blocks of the output in parallel.
code: def loss(self, logits, features):
      assert self._hparams.block_size > 0
  
      def shift_left_4d(x, k):
          return tf.pad(x, [[0, 0], [0, k], [0, 0], [0, 0]])[
              :, k:, :, :
          ]
  
      targets = features["targets"]
      assert len(targets.shape) == 4
  
      targets = tf.concat(
          [
              shift_left_4d(targets, i)
              for i in range(self._hparams.block_size)
          ],
          axis=2,
      )
  
      if (
          self._hparams.mode == tf_estimator.ModeKeys.TRAIN
          or self._hparams.mode == tf_estimator.ModeKeys.EVAL
      ):
          assert "block_index" in features
          k = features["block_index"]
          targets = targets[:, :, k : k + 1, :]
  
      features["targets"] = targets
  
      loss = super(TransformerBlockParallel, self).loss(
          logits, features
      )
  
      if self._hparams.mode == tf_estimator.ModeKeys.TRAIN:
          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
  from tensor2tensor.data_generators import text_encoder
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 3 for calling the TransformerBlockParallel obj's underlying member method code with example usage: obj.loss(logits, features) and returns: loss, tf