def max_pool(x_input, pool_size):
    """max_pool downsamples a feature map by 2X."""
    return tf.nn.max_pool(x_input, ksize=[1, pool_size, pool_size, 1],
                          strides=[1, pool_size, pool_size, 1], padding='SAME')