Example showing PixelShuffle ONNX export doesn't use same ordering as PyTorch

test_pixelshuffle.py

import onnx
import torch
from caffe2.python.onnx import backend
from torch import nn
from torch.autograd import Variable
from torch.nn import PixelShuffle

class MyPixelShuffle(nn.Module):
    """
    Alternative version of pixel shuffle that was copied from the C++ version
    """
    def __init__(self, upscale_factor):
        super(MyPixelShuffle, self).__init__()
        self._upscale_factor = upscale_factor
        self._upscale_factor_squared = upscale_factor ** 2

    def forward(self, x):
        b = x.shape[0]
        c = x.shape[1]
        h = x.shape[2]
        w = x.shape[3]
        oc = c / self._upscale_factor_squared
        oh = h * self._upscale_factor
        ow = w * self._upscale_factor
        input_reshaped = x.reshape(b, oc, self._upscale_factor, self._upscale_factor, h, w)
        permuted = input_reshaped.permute(0, 1, 4, 2, 5, 3)
        return permuted.reshape(b, oc, oh, ow)


not_working = PixelShuffle(2)
working = MyPixelShuffle(2)
shape = (1, 4, 2, 2)
input = Variable(torch.randn(*shape), requires_grad=True)

expected_out = not_working(input).detach().numpy()
print('Expected first row\t\t{}'.format(expected_out[0, 0, 0]))

my_out = working(input).detach().numpy()
print('MyPixelShuffle first row\t{}'.format(my_out[0, 0, 0]))

dummy_input = torch.randn(*shape)
torch.onnx.export(not_working, dummy_input, 'pixel_shuffle.onnx')
model = onnx.load("pixel_shuffle.onnx")
rep = backend.prepare(model)  
outputs = rep.run(input.detach().numpy())[0]
print('Builtin shuffle ONNX first row\t{}'.format(outputs[0, 0, 0]))

torch.onnx.export(working, dummy_input, 'pixel_shuffle_working.onnx')
model = onnx.load("pixel_shuffle_working.onnx")
rep = backend.prepare(model) 
outputs = rep.run(input.detach().numpy())[0]
print('MyPixelShuffle ONNX first row\t{}'.format(outputs[0, 0, 0]))