From 8237f15cfec852f278c28cf542a00da720553626 Mon Sep 17 00:00:00 2001
From: vandanavk <vandana268@gmail.com>
Date: Mon, 5 Nov 2018 08:50:29 -0800
Subject: [PATCH] ONNX import/export: Size

---
 python/mxnet/contrib/onnx/mx2onnx/_op_translations.py | 7 +++++++
 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, 17 insertions(+), 3 deletions(-)

diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
index 7e84cea5ead5..988092910778 100644
--- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
+++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -1545,3 +1545,10 @@ def convert_sum(node, **kwargs):
             name=name
         )
     return [node]
+
+@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 678435d92357..872fe95b5eb5 100644
--- a/tests/python-pytest/onnx/export/onnx_backend_test.py
+++ b/tests/python-pytest/onnx/export/onnx_backend_test.py
@@ -94,7 +94,8 @@
     'test_operator_permute2',
     'test_clip'
     'test_cast',
-    'test_depthtospace'
+    'test_depthtospace',
+    '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 7f34247c94e2..c407f91c7a64 100644
--- a/tests/python-pytest/onnx/import/test_cases.py
+++ b/tests/python-pytest/onnx/import/test_cases.py
@@ -86,7 +86,8 @@
     'test_operator_maxpool',
     'test_operator_params',
     'test_operator_permute2',
-    'test_depthtospace'
+    'test_depthtospace',
+    'test_size'
     ]
 
 BASIC_MODEL_TESTS = [