import tensorflow as tf

class CommonImageAttentionTest(parameterized.TestCase, tf.test.TestCase):

  @parameterized.parameters(
      (common_image_attention.DistributionType.DMOL, 5, 50),
      (common_image_attention.DistributionType.CAT, None, 256),
  )
  def testPostProcessImageTrainMode(self, likelihood, num_mixtures, depth):
    batch = 1
    rows = 8
    cols = 24
    hparams = tf.contrib.training.HParams(
        hidden_size=2,
        likelihood=likelihood,
        mode=tf.estimator.ModeKeys.TRAIN,
        num_mixtures=num_mixtures,
    )
    inputs = tf.random_uniform([batch, rows, cols, hparams.hidden_size],
                               minval=-1., maxval=1.)
    outputs = common_image_attention.postprocess_image(
        inputs, rows, cols, hparams)
