class_docstr: Attention net.  See file docstring.
code: def body(self, features):
      # Remove dropout if not training
      hparams = self._hparams
      ps_devices = self._ps_devices
      assert hparams.num_model_shards % len(ps_devices) == 0
      shards_per_device = hparams.num_model_shards // len(ps_devices)
      model_devices = [
          ps_devices[i // shards_per_device]
          for i in range(hparams.num_model_shards)
      ]
      print("model_devices = %s" % model_devices)
      mp = expert_utils.Parallelism(model_devices, reuse=False)
      vocab_size = self._problem_hparams.vocabulary[
          "targets"
      ].vocab_size
      # squeeze out channels, heights
      targets = features["targets_raw"]
      targets = tf.squeeze(targets, 3)
      targets = tf.squeeze(targets, 2)
      shifted_targets = common_layers.shift_right_2d(targets)
      # Bypass the symbol modality and use a different embedding on each shard.
      decoder_input = mp(
          common_layers.embedding,
          shifted_targets,
          vocab_size,
          hparams.hidden_size,
          multiplier=hparams.hidden_size**0.5,
          symbol_dropout_rate=hparams.symbol_dropout,
      )
      decoder_self_attention_bias = mp(
          common_attention.attention_bias_lower_triangle,
          tf.shape(targets)[1],
      )
      if "targets_segmentation" in features:
          # "Packed" dataset - keep the examples from seeing each other.
          targets_segmentation = features["targets_segmentation"]
          targets_position = features["targets_position"]
          decoder_self_attention_bias = mp(
              tf.add,
              decoder_self_attention_bias,
              mp(
                  common_attention.attention_bias_same_segment,
                  targets_segmentation,
                  targets_segmentation,
              ),
          )
      else:
          targets_position = None
  
      if hparams.pos == "timing":
          if targets_position is None:
              decoder_input = mp(
                  common_attention.add_timing_signal_1d, decoder_input
              )
          else:
              decoder_input = mp(
                  common_attention.add_timing_signal_1d_given_position,
                  decoder_input,
                  targets_position,
              )
  
      decoder_input = mp(
          tf.nn.dropout,
          decoder_input,
          1.0 - hparams.layer_prepostprocess_dropout,
      )
      decoder_output, extra_loss = _super_stack(
          decoder_input, decoder_self_attention_bias, hparams, mp
      )
      # Bypass the symbol modality and compute logits directly.
      # We compute a different set of logits on each shard, and sum them.
      logits = mp(
          tf.layers.dense, decoder_output, vocab_size, name="logits"
      )
      logits = expert_utils.all_reduce_ring(logits, mp)
      logits = mp(tf.multiply, logits, mp.n**-0.5)
      # We now have identical logits on all shards.
      # Shard 0 gets returned to the estimator.
      logits_shard_0 = logits[0]
      logits_shard_0 = tf.expand_dims(logits_shard_0, 2)
      logits_shard_0 = tf.expand_dims(logits_shard_0, 3)
      # On each device, we compute the loss for a part of the batch.
      # This is faster than computing the whole loss on one shard.
      mp, logits = expert_utils.reduce_by_device(
          mp, logits, lambda l: l[0]
      )
  
      def _loss_for_shard(logits, targets, shard):
          if mp.n > 1:
              logits = common_layers.approximate_split(logits, mp.n, 0)[
                  shard
              ]
              targets = common_layers.approximate_split(
                  targets, mp.n, 0
              )[shard]
          return common_layers.padded_cross_entropy(
              logits, targets, hparams.label_smoothing
          )
  
      num, denom = mp(_loss_for_shard, logits, targets, range(mp.n))
      # override training loss so that it is not computed externally.
      losses = {"training": tf.add_n(num) / tf.add_n(denom)}
      if extra_loss is not None:
          losses["extra"] = extra_loss
      return logits_shard_0, losses
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from six.moves import range
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import modalities
  from tensor2tensor.utils import diet
  from tensor2tensor.utils import expert_utils
  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 1 for calling the SuperLM obj's underlying member method code with example usage: obj.body(features) and returns: logits_shard_0,  losses, common_layers