From bbab1c704343e0a3dfb4ae05c0e7cb19b0f6473d Mon Sep 17 00:00:00 2001
From: zha0q1 <zhaoqizhu96@gmail.com>
Date: Fri, 15 Jan 2021 21:34:03 +0000
Subject: [PATCH] _maximum_scalar

---
 .../contrib/onnx/mx2onnx/_op_translations.py  | 40 +++++++++++++++----
 tests/python-pytest/onnx/test_operators.py    |  8 ++++
 2 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
index ecfb0320e593..84a8c8674d78 100644
--- a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
+++ b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -179,19 +179,21 @@ def create_const_node(input_name, value, kwargs):
     initializer.append(tensor_node)
     return value_node
 
-def create_tensor(shape_list, shape_name, initializer, dtype='int64'):
+def create_tensor(tensor_list, tensor_name, initializer, dtype='int64'):
     """Helper function to create a tensor value node and a
     initializer tensor node with constant value."""
-    shape_np = np.array(shape_list, dtype=dtype)
-    data_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[shape_np.dtype]
-    dims = np.shape(shape_np)
-    tensor_node = onnx.helper.make_tensor_value_info(shape_name, data_type, dims)
+    tensor_np = np.array(tensor_list, dtype=dtype)
+    data_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[tensor_np.dtype]
+    dims = np.shape(tensor_np)
+    tensor_node = onnx.helper.make_tensor_value_info(tensor_name, data_type, dims)
+    if dtype == np.float16:
+        tensor_list = tensor_np.view(dtype=np.uint16).flatten().tolist()
     initializer.append(
         onnx.helper.make_tensor(
-            name=shape_name,
+            name=tensor_name,
             data_type=data_type,
             dims=dims,
-            vals=shape_list,
+            vals=tensor_list,
             raw=False
         )
     )
@@ -2946,3 +2948,27 @@ def convert_where(node, **kwargs):
         make_node("Where", [name+"_bool", input_nodes[1], input_nodes[2]], [name], name=name)
     ]
     return nodes
+
+
+@mx_op.register('_maximum_scalar')
+def convert_maximum_scalar(node, **kwargs):
+    """Map MXNet's _maximum_scalar
+    """
+    from onnx.helper import make_node
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    input_type = int(kwargs['in_type'])
+    dtype = onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[input_type]
+
+    scalar = None
+    if 'float' in str(dtype):
+        scalar = float(attrs.get('scalar', '0'))
+    else:
+        scalar = int(attrs.get('scalar', '0'))
+
+    nodes = [
+        create_tensor([scalar], name+'_scalar', kwargs['initializer'], dtype=dtype),
+        make_node('Max', [input_nodes[0], name+'_scalar'], [name], name=name)
+        ]
+
+    return nodes
diff --git a/tests/python-pytest/onnx/test_operators.py b/tests/python-pytest/onnx/test_operators.py
index f9a6abcaf14e..9425e3c5e216 100644
--- a/tests/python-pytest/onnx/test_operators.py
+++ b/tests/python-pytest/onnx/test_operators.py
@@ -408,3 +408,11 @@ def test_onnx_export_where(tmp_path, dtype, shape):
     y = mx.nd.ones(shape, dtype=dtype)
     cond = mx.nd.random.randint(low=0, high=1, shape=shape, dtype='int32')
     op_export_test('where', M, [cond, x, y], tmp_path)
+
+
+@pytest.mark.parametrize('dtype', ['float16', 'float32', 'float64', 'int32', 'int64'])
+@pytest.mark.parametrize('shape', [(3, 4, 5), (1, 4, 1, 7)])
+def test_onnx_maximum_scalar(tmp_path, dtype, shape):
+    x = mx.random.uniform(0, 10, shape).astype(dtype)
+    M = def_model('maximum', right=5)
+    op_export_test('_maximum_scalar', M, [x], tmp_path)