Chainer implementation of pixel shuffle used in sub-pixel convolutions

gistfile1.py

import numpy as np
from chainer import Variable
from chainer import functions as F

# Pixel shuffle used in sub-pixel convolutions, in Chainer
# https://arxiv.org/pdf/1609.05158v2.pdf
#
# Example:
# Scaling factor: 3
# In shape: (1, 9, 4, 4)
# Out shape: (1, 1, 12, 12)
#
# `np` may be replaced with `cupy` (after `import cupy`) to perform
# the same computations on the GPU

upscale_factor = 3
x = np.empty((1, 9, 4, 4), dtype=np.float32)

n, c, w, h = x.shape

# Set all values in each feature map to its feature map index
for i in range(c):
    x[0, i] = i

c_out = c // upscale_factor ** 2
w_out = w * upscale_factor
h_out = h * upscale_factor

x = Variable(x)
x = F.reshape(x, (n, c_out, upscale_factor, upscale_factor, w, h))
x = F.transpose(x, (0, 1, 4, 2, 5, 3))
x = F.reshape(x, (n, c_out, w_out, h_out))

assert(x.shape == (1, 1, 12, 12))