code: def test_transformer_aux_body(self):
      batch_size = 3
      input_length = 5
      target_length = 16
      vocab_size = 9
      hparams = transformer_aux.transformer_aux_tiny()
      hparams.shift_values = "-5,1,2,3"
      p_hparams = problem_hparams.test_problem_hparams(
          vocab_size, vocab_size, hparams
      )
      hparams.problem_hparams = p_hparams
      inputs = np.random.randint(
          vocab_size, size=(batch_size, input_length, 1, 1)
      )
      targets = np.random.randint(
          vocab_size, size=(batch_size, target_length, 1, 1)
      )
      features = {
          "inputs": tf.constant(inputs, dtype=tf.int32),
          "targets": tf.constant(targets, dtype=tf.int32),
          "target_space_id": tf.constant(1, dtype=tf.int32),
      }
      tf.train.create_global_step()
      model = transformer_aux.TransformerAux(
          hparams, tf_estimator.ModeKeys.TRAIN, p_hparams
      )
      logits, losses = model(features)
  
      self.assertIn("training", losses)
      self.assertIn("auxiliary", losses)
  
      with self.test_session() as session:
          session.run(tf.global_variables_initializer())
          logits_val = session.run(logits)
          self.assertEqual(
              logits_val.shape,
              (batch_size, target_length, 1, 1, vocab_size),
          )
  
file_path: tensor2tensor/tensor2tensor/models/research/transformer_aux_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  from tensor2tensor.data_generators import problem_hparams
  from tensor2tensor.models.research import transformer_aux
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 2 for calling the TransformerAuxTest obj's underlying member method code with example usage: obj.test_transformer_aux_body() without return types