import tensorflow as tf

    channels = 3
    rows = height
    if likelihood == common_image_attention.DistributionType.CAT:
      cols = channels * width
    else:
      cols = width
    hparams = tf.contrib.training.HParams(
        hidden_size=2,
        likelihood=likelihood,
        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., maxval=1.)
    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))

if __name__ == "__main__":
  tf.test.main()
