From a2e0cb236b89ec9313114b5a84622ada7f58f8d4 Mon Sep 17 00:00:00 2001 From: vandanavk Date: Thu, 25 Oct 2018 13:58:59 -0700 Subject: [PATCH] ONNX export: Shape operator --- python/mxnet/contrib/onnx/mx2onnx/_op_translations.py | 7 +++++++ tests/python-pytest/onnx/export/onnx_backend_test.py | 3 ++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py index 6605e604fcd3..39bf25f24da3 100644 --- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py +++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py @@ -1562,3 +1562,10 @@ def convert_sum(node, **kwargs): name=name ) return [node] + +@mx_op.register("shape_array") +def convert_shape(node, **kwargs): + """Map MXNet's shape_array operator attributes to onnx's Shape operator + and return the created node. + """ + return create_basic_op_node('Shape', node, kwargs) diff --git a/tests/python-pytest/onnx/export/onnx_backend_test.py b/tests/python-pytest/onnx/export/onnx_backend_test.py index fc30dbdb06da..74c28729c8e0 100644 --- a/tests/python-pytest/onnx/export/onnx_backend_test.py +++ b/tests/python-pytest/onnx/export/onnx_backend_test.py @@ -95,7 +95,8 @@ 'test_clip' 'test_cast', 'test_depthtospace', - 'test_instancenorm' + 'test_instancenorm', + 'test_shape' ] BASIC_MODEL_TESTS = [