code: def testCreateOutputTrainMode(self, likelihood, num_mixtures, depth):
      batch = 1
      height = 8
      width = 8
      channels = 3
      rows = height
      if likelihood == common_image_attention.DistributionType.CAT:
          cols = channels * width
      else:
          cols = width
      hparams = hparam.HParams(
          hidden_size=2,
          likelihood=likelihood,
          num_channels=channels,
          mode=tf_estimator.ModeKeys.TRAIN,
          num_mixtures=num_mixtures,
      )
      decoder_output = tf.random_normal(
          [batch, rows, cols, hparams.hidden_size]
      )
      targets = tf.random_uniform(
          [batch, height, width, channels], minval=-1.0, maxval=1.0
      )
      output = common_image_attention.create_output(
          decoder_output, rows, cols, targets, hparams
      )
      if (
          hparams.likelihood
          == common_image_attention.DistributionType.CAT
      ):
          self.assertEqual(
              output.shape, (batch, height, width, channels, depth)
          )
      else:
          self.assertEqual(output.shape, (batch, height, width, depth))
  
file_path: tensor2tensor/tensor2tensor/layers/common_image_attention_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.utils import hparam
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 3 for calling the CommonImageAttentionTest obj's underlying member method code with example usage: obj.testCreateOutputTrainMode(likelihood, num_mixtures, depth) without return types