dense_split.py

import mxnet as mx
x = mx.np.random.normal(0,1,(1,1,512))
table = mx.np.random.normal(0,1,(128,30))
extra_table = mx.np.random.normal(0,1,(384,30))
bias = mx.np.random.normal(0,1,(30,))

y = mx.np.concatenate([table, extra_table], axis=0)
res1 = mx.np.dot(x,y) + bias
a = mx.np.dot(x[:,:,:128], table)
b = mx.np.dot(x[:,:,128:], extra_table)
res2 = a+b+bias

from numpy.testing import assert_allclose
assert_allclose(res1.asnumpy(), res2.asnumpy(), 1e-5, 1e-5)