[v1.x] Add onnx export function for log2 operator

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

@@ -3780,3 +3780,27 @@ def convert_batch_dot(node, **kwargs):
         ]
 
     return nodes
+
+
+@mx_op.register("log2")
+def convert_log2(node, **kwargs):
+    """Map MXNet's log2 operator attributes to onnx's operator.
+    """
+    from onnx.helper import make_node, make_tensor
+    name, input_nodes, _ = get_inputs(node, kwargs)
+
+    input_type = kwargs["in_type"]
+    dtype = onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[input_type]
+
+    ln2 = np.array([0.693147180559945309], dtype=dtype)
+    if dtype == 'float16':
+        ln2 = ln2.view(dtype=np.uint16)
+    ln2v = make_tensor(name+'_ln2', input_type, [1], ln2)
+
+    nodes = [
+        make_node('Log', [input_nodes[0]], [name+'_log']),
+        make_node('Constant', [], [name+'_ln2'], value=ln2v),
+        make_node('Div', [name+'_log', name+'_ln2'], [name], name=name)
+    ]
+
+    return nodes

tests/python-pytest/onnx/test_operators.py

@@ -69,9 +69,9 @@ def onnx_rt(onnx_file, inputs):
         pred_nat = pred_nat[0]
     if isinstance(pred_nat, list):
         for i in range(len(pred_nat)):
-            assert_almost_equal(pred_nat[i], pred_onx[i])
+            assert_almost_equal(pred_nat[i], pred_onx[i], equal_nan=True)
     else:
-        assert_almost_equal(pred_nat, pred_onx[0])
+        assert_almost_equal(pred_nat, pred_onx[0], equal_nan=True)
 
 
 def test_onnx_export_abs(tmp_path):
@@ -740,6 +740,13 @@ def test_onnx_export_batch_dot(tmp_path, dtype, transpose_a, transpose_b):
     op_export_test('batch_dot2', M2, [x2, y2], tmp_path)
 
 
+@pytest.mark.parametrize('dtype', ['float16', 'float32'])
+def test_onnx_export_log2(tmp_path, dtype):
+    x = mx.random.normal(0, 10, (2, 3, 4, 5)).astype(dtype)
+    M = def_model('log2')
+    op_export_test('log2', M, [x], tmp_path)
+
+
 @pytest.mark.parametrize('dtype', ['int32', 'int64', 'float16', 'float32', 'float64'])
 @pytest.mark.parametrize('axis', [None, 1, [1,2], -1])
 def test_onnx_export_sum(tmp_path, dtype, axis):