import tensorflow as tf

    batch = 1
    rows = 8
    cols = 24
    block_length = 4
    block_width = 2
    hparams = tf.contrib.training.HParams(
        block_raster_scan=True,
        hidden_size=2,
        likelihood=likelihood,
        mode=tf.estimator.ModeKeys.PREDICT,
        num_mixtures=num_mixtures,
        query_shape=[block_length, block_width],
    )
    inputs = tf.random_uniform([batch, rows, cols, hparams.hidden_size],
                               minval=-1., maxval=1.)
    outputs = common_image_attention.postprocess_image(
        inputs, rows, cols, hparams)
    num_blocks_rows = rows // block_length
    num_blocks_cols = cols // block_width
    self.assertEqual(outputs.shape,
                     (batch, num_blocks_rows, num_blocks_cols,
                      block_length, block_width, depth))

  @parameterized.parameters(
      (common_image_attention.DistributionType.DMOL, 5, 50),
      (common_image_attention.DistributionType.CAT, None, 256),
