Description

When using MKLDNN, concatenating sometimes fails and results in error. Please see the code.

Error Message

MXNetError: The size of NDArray doesn't match the requested MKLDNN memory desc. MKLDNN memory requests for 41943040 bytes, but got 34603008 bytes from NDArray

To Reproduce

x = mx.sym.Variable('x')
w = mx.sym.Variable('w')

c = mx.sym.Convolution(data=x, weight=w, num_filter=32, kernel=(3, 3), pad=(1, 1), no_bias=True)
n = mx.sym.concat(c, x, dim=1)

context = mx.cpu()
ex = n.simple_bind(context, x=(1, 1, 512, 512))

data = mx.nd.ones(ex.arg_arrays[0].shape, ctx=context)
weight = mx.nd.ones(ex.arg_arrays[1].shape, ctx=context)

data.copyto(ex.arg_arrays[0])
weight.copyto(ex.arg_arrays[1])

res = ex.forward()
print(res)

Steps to reproduce

Simply run the provided script

What have you tried to solve it?

If variable x is first tiled to have at least 8 channels (dimension 1), then it works:
c = mx.sym.Convolution(data=x, weight=w, num_filter=32, kernel=(3, 3), pad=(1, 1), no_bias=True)
x = mx.sym.tile(x, (1, 8, 1, 1))
n = mx.sym.concat(c, x, dim=1)

Another way is to reverse the order of concatenation:
n = mx.sym.concat(x, c, dim=1)

Environment

MXNet 1.8.0.rc2, Windows 10, build with MKLDNN