add aten::pixel_shuffle implementation (#6328) #6468
Conversation
|
thanks, please add a test. |
Thanks! The test has been added! |
| out_shape = [b, oc, oh, ow] | ||
|
|
||
| data = _op.transform.reshape(data, new_shape) | ||
| axes = [0, 1, 4, 2, 5, 3] |
There was a problem hiding this comment.
Can you tell me what are these hardcoded axes?
There was a problem hiding this comment.
The original shape order is [b, oc, upscale_factor, upscale_factor, h, w], the hardcoded axes transpose the order to [b, oc, h, upscale_factor, w, upscale_factor] for further reshape
There was a problem hiding this comment.
ok, please add your explanation above as a comment (so that people won't get confused).
There was a problem hiding this comment.
The comment has been added.
|
@wjliu1998 lint is complaining |
python/tvm/relay/frontend/pytorch.py
Outdated
| raise AssertionError(msg) | ||
|
|
||
| if isinstance(data, tvm.runtime.NDArray): | ||
| ndims = len(_infer_shape(data)) |
There was a problem hiding this comment.
It seems you copy pasted the above code from _transpose(). That one is a legacy code we don't want to repeat. Please leave only what is necessary.
There was a problem hiding this comment.
Thanks! The code has been deleted
There was a problem hiding this comment.
I mean, you have three if isinstance(...) right? You shouldn't need isintance, just get ndims directly (figure out which code path actually hit in your test). I guess all you need is just ndims = data.
There was a problem hiding this comment.
Thank you for pointing out that. It seems ndims = len(_infer_shape(data, prelude.mod)) is enough
|
Thanks @wjliu1998 |
Following the implementation of aten::pixel_shuffle in pytorch. The implementation of pixel_shuffle has been added in pytorch frontend