Keras only 2x pixel-shuffle / sub_pixel_conv

sub_pixel_conv.py

def sub_pixel_conv(inputs, height, width, out_channels): #keras version of 2x pixel shuffle
    x = inputs
    x = SeparableConv2D(out_channels * 4, kernel_size = 9, depth_multiplier = 1, activation = 'tanh', padding = 'same')(x)
    x = Reshape((height, width, out_channels, 2, 2))(x)
    x = Permute((3, 2, 4, 1, 5))(x)
    x = Reshape((out_channels, height * 2, height * 2))(x)
    x = Permute((2, 3, 1))(x)
    return x

like tf.depth_to_space