import tensorflow as tf

        return input_
    shape = input_.get_shape().as_list()
    batch_size = shape[0]
    height = shape[1]
    width = shape[2]
    channels = shape[3]
    if channels % 4 != 0:
        raise ValueError("Number of channels not divisible by 4.")
    res = tf.reshape(input_, [batch_size, height, width, channels // 4, 2, 2])
    res = tf.transpose(res, [0, 1, 4, 2, 5, 3])
    res = tf.reshape(res, [batch_size, 2 * height, 2 * width, channels // 4])

    return res


# batch norm
def batch_norm(input_,
               dim,
