def testDepthInterleavedDepth3(self):
   x_np = [[[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]]]]
   block_size = 2
   with self.test_session(use_gpu=False):
     x_tf = tf.depth_to_space(x_np, block_size)
     self.assertAllEqual(x_tf.eval(), [[[[1, 2, 3], [4, 5, 6]],
                                        [[7, 8, 9], [10, 11, 12]]]])

def SubpixelConv2D(*args, **kwargs):
    kwargs['output_dim'] = 4*kwargs['output_dim']
    output = lib.ops.conv2d.Conv2D(*args, **kwargs)
    output = tf.transpose(output, [0,2,3,1])
    output = tf.depth_to_space(output, 2)
    output = tf.transpose(output, [0,3,1,2])
    return output

 def testBlockSize0(self):
   x_np = [[[[1], [2]],
            [[3], [4]]]]
   block_size = 0
   with self.assertRaises(ValueError):
     out_tf = tf.depth_to_space(x_np, block_size)
     out_tf.eval()

def depth_to_space(input, scale, data_format=None):
    ''' Uses phase shift algorithm to convert channels/depth for spatial resolution '''
    data_format = 'NHWC'

    data_format = data_format.lower()
    out = tf.depth_to_space(input, scale, data_format=data_format)
    return out

 def testDepthInterleaved(self):
   x_np = [[[[1, 10, 2, 20, 3, 30, 4, 40]]]]
   block_size = 2
   with self.test_session(use_gpu=False):
     x_tf = tf.depth_to_space(x_np, block_size)
     self.assertAllEqual(x_tf.eval(), [[[[1, 10], [2, 20]],
                                        [[3, 30], [4, 40]]]])

 def deconv2d(cur, i):
   thicker = conv(
       cur,
       output_filters * 4, (1, 1),
       padding="SAME",
       activation=tf.nn.relu,
       name="deconv2d" + str(i))
   return tf.depth_to_space(thicker, 2)

def UpsampleConv(name, input_dim, output_dim, filter_size, inputs, he_init=True, biases=True):
    output = inputs
    output = tf.concat([output, output, output, output], axis=1)
    output = tf.transpose(output, [0,2,3,1])
    output = tf.depth_to_space(output, 2)
    output = tf.transpose(output, [0,3,1,2])
    output = lib.ops.conv2d.Conv2D(name, input_dim, output_dim, filter_size, output, he_init=he_init, biases=biases)
    return output

 def testBlockSizeOne(self):
   x_np = [[[[1, 1, 1, 1],
             [2, 2, 2, 2]],
            [[3, 3, 3, 3],
             [4, 4, 4, 4]]]]
   block_size = 1
   with self.assertRaises(ValueError):
     out_tf = tf.depth_to_space(x_np, block_size)
     out_tf.eval()

 def testBlockSize4FlatInput(self):
   x_np = [[[[1, 2, 5, 6, 3, 4, 7, 8, 9, 10, 13, 14, 11, 12, 15, 16]]]]
   block_size = 4
   with self.test_session(use_gpu=False):
     x_tf = tf.depth_to_space(x_np, block_size)
     self.assertAllEqual(x_tf.eval(), [[[[1], [2], [5], [6]],
                                        [[3], [4], [7], [8]],
                                        [[9], [10], [13], [14]],
                                        [[11], [12], [15], [16]]]])

 def depth_to_space(cls, ipt, scale, data_format=None):
     """ Uses phase shift algorithm to convert channels/depth
         for spatial resolution """
     if data_format is None:
         data_format = K.image_data_format()
     data_format = data_format.lower()
     ipt = cls._preprocess_conv2d_input(ipt, data_format)
     out = tf.depth_to_space(ipt, scale)
     out = cls._postprocess_conv2d_output(out, data_format)
     return out

def phase_shift(x, upsampling_factor=2, data_format="NCHW", name="PhaseShift"):

    if data_format == "NCHW":
        x = tf.transpose(x, [0,2,3,1])

    x = tf.depth_to_space(x, upsampling_factor, name=name)

    if data_format == "NCHW":
        x = tf.transpose(x, [0,3,1,2])

    return x

 def testDepthToSpaceTranspose(self):
   x = np.arange(20 * 5 * 8 * 7, dtype=np.float32).reshape([20, 5, 8, 7])
   block_size = 2
   crops = np.zeros((2, 2), dtype=np.int32)
   y1 = tf.batch_to_space(x, crops, block_size=block_size)
   y2 = tf.transpose(
       tf.depth_to_space(
           tf.transpose(x, [3, 1, 2, 0]), block_size=block_size),
       [3, 1, 2, 0])
   with self.test_session():
     self.assertAllEqual(y1.eval(), y2.eval())

 def testBlockSizeTooLarge(self):
   x_np = [[[[1, 2, 3, 4],
             [5, 6, 7, 8]],
            [[9, 10, 11, 12],
             [13, 14, 15, 16]]]]
   block_size = 4
   # Raise an exception, since th depth is only 4 and needs to be
   # divisible by 16.
   with self.assertRaises(IndexError):
     out_tf = tf.depth_to_space(x_np, block_size)
     out_tf.eval()

def decompress_step(source, hparams, first_relu, is_2d, name):
  """Decompression function."""
  with tf.variable_scope(name):
    shape = common_layers.shape_list(source)
    multiplier = 4 if is_2d else 2
    kernel = (1, 1) if is_2d else (1, 1)
    thicker = common_layers.conv_block(
        source, hparams.hidden_size * multiplier, [((1, 1), kernel)],
        first_relu=first_relu, name="decompress_conv")
    if is_2d:
      return tf.depth_to_space(thicker, 2)
    return tf.reshape(thicker, [shape[0], shape[1] * 2, 1, hparams.hidden_size])

 def testDepthInterleavedLarger(self):
   x_np = [[[[1, 10, 2, 20, 3, 30, 4, 40],
             [5, 50, 6, 60, 7, 70, 8, 80]],
            [[9, 90, 10, 100, 11, 110, 12, 120],
             [13, 130, 14, 140, 15, 150, 16, 160]]]]
   block_size = 2
   with self.test_session(use_gpu=False):
     x_tf = tf.depth_to_space(x_np, block_size)
     self.assertAllEqual(x_tf.eval(),
                         [[[[1, 10], [2, 20], [5, 50], [6, 60]],
                           [[3, 30], [4, 40], [7, 70], [8, 80]],
                           [[9, 90], [10, 100], [13, 130], [14, 140]],
                           [[11, 110], [12, 120], [15, 150], [16, 160]]]])

 def testNonSquare(self):
   x_np = [[[[1, 10, 2, 20, 3, 30, 4, 40]],
            [[5, 50, 6, 60, 7, 70, 8, 80]],
            [[9, 90, 10, 100, 11, 110, 12, 120]]]]
   block_size = 2
   with self.test_session(use_gpu=False):
     x_tf = tf.depth_to_space(x_np, block_size)
     self.assertAllEqual(x_tf.eval(), [[[[1, 10], [2, 20]],
                                        [[3, 30], [4, 40]],
                                        [[5, 50], [6, 60]],
                                        [[7, 70], [8, 80]],
                                        [[9, 90], [10, 100]],
                                        [[11, 110], [12, 120]]]])

    def _PS(X, r, n_out_channels):
        if n_out_channels >= 1:
            assert int(X.get_shape()[-1]) == (r**2) * n_out_channels, _err_log

            # bsize, a, b, c = X.get_shape().as_list()
            # bsize = tf.shape(X)[0] # Handling Dimension(None) type for undefined batch dim
            # Xs=tf.split(X,r,3) #b*h*w*r*r
            # Xr=tf.concat(Xs,2) #b*h*(r*w)*r
            # X=tf.reshape(Xr,(bsize,r*a,r*b,n_out_channel)) # b*(r*h)*(r*w)*c

            X = tf.depth_to_space(X, r)
        else:
            logging.info(_err_log)
        return X

  def _checkGrad(self, x, block_size):
    assert 4 == x.ndim
    with self.test_session():
      tf_x = tf.convert_to_tensor(x)
      tf_y = tf.depth_to_space(tf_x, block_size)
      epsilon = 1e-2
      ((x_jacob_t, x_jacob_n)) = tf.test.compute_gradient(
          tf_x,
          x.shape,
          tf_y,
          tf_y.get_shape().as_list(),
          x_init_value=x,
          delta=epsilon)

    self.assertAllClose(x_jacob_t, x_jacob_n, rtol=1e-2, atol=epsilon)

    def layer_conv_dts(self, net, args, options):
        options = hc.Config(options)
        config = self.config
        ops = self.ops

        self.ops.activation_name = options.activation_name
        activation_s = options.activation or self.ops.config_option("activation")
        activation = self.ops.lookup(activation_s)

        stride = options.stride or self.ops.config_option("stride", [1,1])[0]
        stride = int(stride)
        fltr = options.filter or self.ops.config_option("filter", [3,3])
        if type(fltr) == type(""):
            fltr=[int(fltr), int(fltr)]
        depth = int(args[0])

        initializer = None # default to global

        trainable = True
        if options.trainable == 'false':
            trainable = False
        bias = True
        if options.bias == 'false':
            bias=False
        net = ops.conv2d(net, fltr[0], fltr[1], stride, stride, depth*4, initializer=initializer, trainable=trainable, bias=bias)
        s = ops.shape(net)
        net = tf.depth_to_space(net, 2)
        if activation:
            #net = self.layer_regularizer(net)
            net = activation(net)

        avg_pool = options.avg_pool or self.ops.config_option("avg_pool")
        if type(avg_pool) == type(""):
            avg_pool = [int(avg_pool), int(avg_pool)]
        if avg_pool:
            ksize = [1,avg_pool[0], avg_pool[1],1]
            stride = ksize
            net = tf.nn.avg_pool(net, ksize=ksize, strides=stride, padding='SAME')

        return net

def network(input):  # Unet
    conv1 = slim.conv2d(input, 32, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv1_1')
    conv1 = slim.conv2d(conv1, 32, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv1_2')
    pool1 = slim.max_pool2d(conv1, [2, 2], padding='SAME')

    conv2 = slim.conv2d(pool1, 64, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv2_1')
    conv2 = slim.conv2d(conv2, 64, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv2_2')
    pool2 = slim.max_pool2d(conv2, [2, 2], padding='SAME')

    conv3 = slim.conv2d(pool2, 128, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv3_1')
    conv3 = slim.conv2d(conv3, 128, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv3_2')
    pool3 = slim.max_pool2d(conv3, [2, 2], padding='SAME')

    conv4 = slim.conv2d(pool3, 256, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv4_1')
    conv4 = slim.conv2d(conv4, 256, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv4_2')
    pool4 = slim.max_pool2d(conv4, [2, 2], padding='SAME')

    conv5 = slim.conv2d(pool4, 512, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv5_1')
    conv5 = slim.conv2d(conv5, 512, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv5_2')

    up6 = upsample_and_concat(conv5, conv4, 256, 512)
    conv6 = slim.conv2d(up6, 256, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv6_1')
    conv6 = slim.conv2d(conv6, 256, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv6_2')

    up7 = upsample_and_concat(conv6, conv3, 128, 256)
    conv7 = slim.conv2d(up7, 128, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv7_1')
    conv7 = slim.conv2d(conv7, 128, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv7_2')

    up8 = upsample_and_concat(conv7, conv2, 64, 128)
    conv8 = slim.conv2d(up8, 64, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv8_1')
    conv8 = slim.conv2d(conv8, 64, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv8_2')

    up9 = upsample_and_concat(conv8, conv1, 32, 64)
    conv9 = slim.conv2d(up9, 32, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv9_1')
    conv9 = slim.conv2d(conv9, 32, [3, 3], rate=1, activation_fn=lrelu, scope='g_conv9_2')

    conv10 = slim.conv2d(conv9, 27, [1, 1], rate=1, activation_fn=None, scope='g_conv10')
    out = tf.depth_to_space(conv10, 3)
    return out