From e08a4ac640537ce4b6412e4040857a9a2a196509 Mon Sep 17 00:00:00 2001
From: Vandana Kannan <vandanavk@users.noreply.github.com>
Date: Fri, 7 Dec 2018 19:39:47 -0800
Subject: [PATCH] ONNX import/export: Size (#13112)

---
 python/mxnet/contrib/onnx/mx2onnx/_op_translations.py | 8 ++++++++
 python/mxnet/contrib/onnx/onnx2mx/_import_helper.py   | 3 ++-
 python/mxnet/contrib/onnx/onnx2mx/_op_translations.py | 4 ++++
 tests/python-pytest/onnx/export/onnx_backend_test.py  | 3 ++-
 tests/python-pytest/onnx/import/test_cases.py         | 3 ++-
 5 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
index 0f4b448a5416..0d20c76240bd 100644
--- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
+++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -1647,3 +1647,11 @@ def convert_logical_not(node, **kwargs):
     and return the created node.
     """
     return create_basic_op_node('Not', node, kwargs)
+
+
+@mx_op.register("size_array")
+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)
diff --git a/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py b/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py
index f61910f838ea..2ceabaec1dcd 100644
--- a/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py
+++ b/python/mxnet/contrib/onnx/onnx2mx/_import_helper.py
@@ -21,7 +21,7 @@
 from ._op_translations import identity, random_uniform, random_normal
 from ._op_translations import add, subtract, multiply, divide, absolute, negative, add_n
 from ._op_translations import tanh, arccos, arcsin, arctan, _cos, _sin, _tan
-from ._op_translations import softplus, shape, gather, lp_pooling
+from ._op_translations import softplus, shape, gather, lp_pooling, size
 from ._op_translations import ceil, floor, hardsigmoid, global_lppooling
 from ._op_translations import concat
 from ._op_translations import leaky_relu, _elu, _prelu, _selu, softmax, fully_connected
@@ -139,6 +139,7 @@
     'Softplus'          : softplus,
     'Tan'               : _tan,
     'Shape'             : shape,
+    'Size'              : size,
     'Gather'            : gather,
     'HardSigmoid'       : hardsigmoid,
     'LpPool'            : lp_pooling,
diff --git a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py
index 368b98d662b1..702832529314 100644
--- a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py
+++ b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py
@@ -642,6 +642,10 @@ def shape(attrs, inputs, proto_obj):
     """Returns shape of input array."""
     return 'shape_array', attrs, inputs
 
+def size(attrs, inputs, proto_obj):
+    """Returns array containing size of data."""
+    return "size_array", attrs, inputs
+
 def reduce_l2(attrs, inputs, proto_obj):
     """Reduce input tensor by l2 normalization."""
     new_attrs = translation_utils._fix_attribute_names(attrs, {'axes':'axis'})
diff --git a/tests/python-pytest/onnx/export/onnx_backend_test.py b/tests/python-pytest/onnx/export/onnx_backend_test.py
index be9273eb6fac..c9926c4d5e15 100644
--- a/tests/python-pytest/onnx/export/onnx_backend_test.py
+++ b/tests/python-pytest/onnx/export/onnx_backend_test.py
@@ -97,7 +97,8 @@
     'test_depthtospace',
     'test_hardsigmoid',
     'test_instancenorm',
-    'test_shape'
+    'test_shape',
+    'test_size'
     ]
 
 BASIC_MODEL_TESTS = [
diff --git a/tests/python-pytest/onnx/import/test_cases.py b/tests/python-pytest/onnx/import/test_cases.py
index f41fe92352db..e0b26cc49830 100644
--- a/tests/python-pytest/onnx/import/test_cases.py
+++ b/tests/python-pytest/onnx/import/test_cases.py
@@ -85,7 +85,8 @@
     'test_operator_maxpool',
     'test_operator_params',
     'test_operator_permute2',
-    'test_depthtospace'
+    'test_depthtospace',
+    'test_size'
     ]
 
 BASIC_MODEL_TESTS = [