[v1.x] Add more onnx operator export unit tests

python/mxnet/onnx/mx2onnx/_op_translations/_op_translations_opset12.py

@@ -2308,7 +2308,15 @@ def convert_size(node, **kwargs):
     """Map MXNet's size_array operator attributes to onnx's Size operator
     and return the created node.
     """
-    return create_basic_op_node('Size', node, kwargs)
+    from onnx.helper import make_node
+    name, input_nodes, _ = get_inputs(node, kwargs)
+
+    create_tensor([1], name+'_1', kwargs['initializer'])
+    nodes = [
+        make_node('Size', [input_nodes[0]], [name+'_size']),
+        make_node('Reshape', [name+'_size', name+'_1'], [name], name=name)
+    ]
+    return nodes
 
 
 @mx_op.register("log_softmax")

tests/python-pytest/onnx/test_operators.py

@@ -1406,4 +1406,42 @@ def test_onnx_export_ufunc(tmp_path, dtype, shape, op_name):
 def test_onnx_export_squeeze(tmp_path, dtype, shape_axis):
     x = mx.nd.random.uniform(1, 100, shape=shape_axis[0]).astype(dtype)
     M = def_model('squeeze', axis=shape_axis[1])
-    op_export_test('squeeze', M, [x], tmp_path)
+    op_export_test('squeeze', M, [x], tmp_path)
+
+
+@pytest.mark.parametrize("dtype", ["float16", "float32", "float64", "int32", "int64"])
+@pytest.mark.parametrize("shape", [(1,2,3), (1,10)])
+@pytest.mark.parametrize("axis", [None, 0, 1])
+def test_onnx_export_rnn_param_concat(tmp_path, dtype, shape, axis):
+    kwargs = {}
+    if axis is not None:
+        kwargs['dim'] = axis
+    M = def_model('_internal._rnn_param_concat', **kwargs)
+    x = mx.nd.random.uniform(-1, 1, shape).astype(dtype)
+    y = mx.nd.random.uniform(-1, 1, shape).astype(dtype)
+    op_export_test('_internal._rnn_param_concat', M, [x, y], tmp_path)
+
+
+@pytest.mark.parametrize("dtype", ["float16", "float32", "float64", "int32", "int64"])
+@pytest.mark.parametrize("shape", [(10,), (1,2,3), (4,5,6)])
+def test_onnx_export_size_array(tmp_path, dtype, shape):
+    M = def_model('size_array')
+    x = mx.nd.random.uniform(-1, 1, shape).astype(dtype)
+    op_export_test('size_array', M, [x], tmp_path)
+
+
+@pytest.mark.parametrize("dtype", ["float16", "float32"])
+@pytest.mark.parametrize("shape", [(1,5), (2,10), (4,5)])
+@pytest.mark.parametrize("sample_shape", [(1), (2)])
+def test_onnx_export_sample_multinomial(tmp_path, dtype, shape, sample_shape):
+    kwargs = {}
+    if sample_shape is not None:
+        kwargs['shape'] = sample_shape
+    M = def_model('sample_multinomial', **kwargs)
+    a = mx.nd.random.uniform(0, 1, shape).astype(dtype)
+    x = a/a.sum(axis=1, keepdims=1)
+    def rand_check(out):
+        return np.zeros_like(out)
+    def rand_check_nd(out):
+        return rand_check(out.asnumpy())
+    op_export_test('sample_multinomial', M, [x], tmp_path, mx_map=rand_check_nd, onnx_map=rand_check)