手动实现pixel_unshuffle时使用 F.conv2d报错

我希望使用megengine手动实现pixel_unshuffle方法
使用F.conv2d(input, kernel, stride=downscale_factor, groups=c)代码中
input.shape=(2, 3, 1980, 2880)
kernel.shape=(48,1,4,4)
downscale_factor=4
c=3

同样的方法在pytorch中有效，输出shape为(2, 48, 495, 720),但是在MegEngine的F.conv2d会爆出错误
## 环境
1.系统环境：
2.MegEngine版本：1.11.1+cu111
3.python版本：3.8

## 复现步骤
1.
2.
3.
pytorch代码：
import torch
import torch.nn as nn
import torch.nn.functional as F

def pixel_unshuffle(input, downscale_factor):
'''
input: batchSize * c * k*w * k*h
kdownscale_factor: k

batchSize * c * k*w * k*h -> batchSize * k*k*c * w * h
'''
c = input.shape[1]

kernel = torch.zeros(size=[downscale_factor * downscale_factor * c,
1, downscale_factor, downscale_factor],
device=input.device)
for y in range(downscale_factor):
for x in range(downscale_factor):
kernel[x + y * downscale_factor::downscale_factor*downscale_factor, 0, y, x] = 1
return F.conv2d(input, kernel, stride=downscale_factor, groups=c)

x = torch.zeros(size=(2, 3, 1980, 2880))

x = pixel_unshuffle(x, 4)
print(x.shape)

megengine代码：
def pixel_unshuffle(input, downscale_factor):
'''
input: (batchSize, c, k*w, k*h)
downscale_factor: k

(batchSize, c, k*w, k*h) -> (batchSize, k*k*c, w, h)
'''
c = input.shape[1]

# 创建一个展开操作符（类似于PyTorch中的卷积核）
kernel = F.zeros((downscale_factor * downscale_factor * c, 1, downscale_factor, downscale_factor), device=input.device)
for y in range(downscale_factor):
for x in range(downscale_factor):
kernel[x + y * downscale_factor::downscale_factor*downscale_factor, 0, y, x] = 1
return F.conv2d(input, kernel, stride=downscale_factor, groups=c)

im = F.zeros((2, 3, 1980, 2880))
isn = pixel_unshuffle(im, 4)
print(isn)

## 请提供完整的日志及报错信息
Traceback (most recent call last):
File "/home/kyoron/.pycharm_helpers/pydev/pydevd.py", line 1500, in _exec
pydev_imports.execfile(file, globals, locals) # execute the script
File "/home/kyoron/.pycharm_helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
exec(compile(contents+"\n", file, 'exec'), glob, loc)
File "/mnt/e/多媒体/桌面/WZ/code/CREStereo_swinT/CREStereo/nets/utils/utils.py", line 65, in
python-BaseException
isn = pixel_unshuffle(im, 4)
File "/mnt/e/多媒体/桌面/WZ/code/CREStereo_swinT/CREStereo/nets/utils/utils.py", line 62, in pixel_unshuffle
return F.conv2d(input, kernel, stride=downscale_factor, groups=c)
File "/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/functional/nn.py", line 265, in conv2d
(output,) = apply(op, inp, weight)
RuntimeError: assertion `filter.ndim == img_ndim + 3 || filter.ndim == img_ndim + 5' failed at ../../../../../../dnn/src/common/convolution.cpp:56: void {anonymous}::make_canonized_filter_meta_nchw_nhwc(size_t, const megdnn::TensorLayout&, const Param&, typename megdnn::ConvolutionBase::CanonizedFilterMeta&) [with Parameter = megdnn::param::Convolution; Param = megdnn::param::Convolution; size_t = long unsigned int; typename megdnn::ConvolutionBase::CanonizedFilterMeta = megdnn::ConvolutionBase::CanonizedFilterMeta]
extra message: bad filter ndim for group convolution: spatial_ndim=2 filter_ndim=4

backtrace:
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/lib/libmegengine_shared.so(_ZN3mgb13MegBrainErrorC1ERKSs+0x4a) [0x7feefd2d76aa]
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/lib/libmegengine_shared.so(+0x2ab2557) [0x7feefd339557]
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/lib/libmegengine_shared.so(_ZN6megdnn12ErrorHandler15on_megdnn_errorERKSs+0x14) [0x7fef00addb34]
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/lib/libmegengine_shared.so(_ZN6megdnn12ErrorHandler15on_megdnn_errorEPKc+0x22) [0x7fef00adf632]
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/lib/libmegengine_shared.so(_ZN6megdnn15__assert_fail__EPKciS1_S1_S1_z+0x190) [0x7fef00b80c50]
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/lib/libmegengine_shared.so(_ZNK6megdnn15ConvolutionBaseINS_5param11ConvolutionEE26make_canonized_filter_metaEmRKNS_12TensorLayoutE+0x12d6) [0x7fef00ae8a46]
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/lib/libmegengine_shared.so(_ZNK6megdnn15ConvolutionBaseINS_5param11ConvolutionEE17deduce_layout_fwdERKNS_12TensorLayoutES6_RS4_+0xa3) [0x7fef00af18c3]
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/lib/libmegengine_shared.so(_ZN6megdnn18ConvolutionForward13deduce_layoutERKNS_12TensorLayoutES3_RS1_+0x1b) [0x7fef00af388b]
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/_imperative_rt.cpython-38-x86_64-linux-gnu.so(+0x33d733) [0x7fef99e1c733]
/home/kyoron/miniconda3/envs/dl/lib/python3.8/site-packages/megengine/core/_imperative_rt.cpython-38-x86_64-linux-gnu.so(+0x2e5bff) [0x7fef99dc4bff]

ERROR conda.cli.main_run:execute(124): `conda run python /home/kyoron/.pycharm_helpers/pydev/pydevd.py --multiprocess --qt-support=auto --client 127.0.0.1 --port 59297 --file /mnt/e/多媒体/桌面/WZ/code/CREStereo_swinT/CREStereo/nets/utils/utils.py` failed. (See above for error)

Start by running the reported MegEngine 1.11.1 reproduction with F.conv2d, using the input and kernel shapes in the issue. Read megengine/functional/nn.py around conv2d and the convolution.cpp assertion for grouped filters. Done means the reported grouped convolution either works with these shapes or the limitation and supported shape requirements are clearly documented.

Written by the indexing model from the issue text.