torch.aten.pixel_shuffle

[newling]

No description provided.

[dan-garvey]

Have you done a torch-mlir op lowering before? This may be a meaty one to start on...

[newling]

I haven't done one yet. I figured it'd just get decomposed as a reshape, a permute, and then another reshape, unless I've misread the description. I haven't really started on it yet though, so happy to change to something else.

Something like

(N, C, H, W) -> reshape -> (N, C/K^2, K, K, H, W) -> permute -> (N, C/K^2, H, K, W, K) -> reshape -> (N, C/K^2, HK, WK)

[harishanand95]

https://github.com/pytorch/pytorch/blob/e1cd220b901c7e1bb37bf822c356ba3f8e6b3134/torch/nn/functional.py#L1921 Here is an older pytorch code for pixel_shuffle that works through torch-mlir, if it helps. I had used this as a workaround before.

[newling]

Thanks @harishanand95 ! That's basically the algorithm that @vivekkhandelwal1 and I had in mind