refactor

Author: zhijxu-MS

tests/test_backend.py

@@ -1128,7 +1128,7 @@ def test_onehot0(self):
             _ = tf.identity(x_, name=_TFOUTPUT)
             self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @unittest.skip("")
+    @unittest.skip("only rank 1 is currently implemented")
     def test_onehot1(self):
         # only rank 1 is currently implemented
         x_val = np.array([[0, 2], [1, -1]], dtype=np.int32)
@@ -1139,12 +1139,14 @@ def test_onehot1(self):
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
     def test_onehot2(self):
-        x_val = np.array([0, 1, 2, 1, 2, 0, 1, 2, 1, 2], dtype=np.int32)
-        depth = 20
-        x = tf.placeholder(tf.int32, x_val.shape, name=_TFINPUT)
-        x_ = tf.one_hot(x, depth, on_value=5.0, axis=-1, off_value=1.0, dtype=tf.float32)
-        _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+        for axis in [-1, 0, 1]:
+            tf.reset_default_graph()
+            x_val = np.array([0, 1, 2, 1, 2, 0, 1, 2, 1, 2], dtype=np.int32)
+            depth = 20
+            x = tf.placeholder(tf.int32, x_val.shape, name=_TFINPUT)
+            x_ = tf.one_hot(x, depth, on_value=5.0, axis=axis, off_value=1.0, dtype=tf.float32)
+            _ = tf.identity(x_, name=_TFOUTPUT)
+            self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
     @check_opset_min_version(9, "onehot")
     def test_onehot3(self):
@@ -1160,16 +1162,17 @@ def test_onehot3(self):
             graph = self._run_test_case([_OUTPUT], {_INPUT: x_val})
             self.assertTrue(len(group_nodes_by_type(graph)["OneHot"]) == 1, "onnx onehot should be used")
         # rank 2
-        for np_dtype, tf_dtype in zip([np.int32, np.int64], [tf.int32, tf.int64]):
-            tf.reset_default_graph()
-            x_val = np.arange(0, 50, dtype=np_dtype).reshape([-1, 10])
-            depth = np.array(20).astype(np.int64)
-            x = tf.placeholder(tf_dtype, x_val.shape, name=_TFINPUT)
-            on_off = np.array([5.6, 1.2]).astype(np_dtype)
-            x_ = tf.one_hot(x, depth, on_value=on_off[0], axis=-1, off_value=on_off[1])
-            _ = tf.identity(x_, name=_TFOUTPUT)
-            graph = self._run_test_case([_OUTPUT], {_INPUT: x_val})
-            self.assertTrue(len(group_nodes_by_type(graph)["OneHot"]) == 1, "onnx onehot should be used")
+        for aixs in [-1, 0, 1, 2]:
+            for np_dtype, tf_dtype in zip([np.int32, np.int64], [tf.int32, tf.int64]):
+                tf.reset_default_graph()
+                x_val = np.arange(0, 50, dtype=np_dtype).reshape([-1, 10])
+                depth = np.array(20).astype(np.int64)
+                x = tf.placeholder(tf_dtype, x_val.shape, name=_TFINPUT)
+                on_off = np.array([5.6, 1.2]).astype(np_dtype)
+                x_ = tf.one_hot(x, depth, on_value=on_off[0], axis=aixs, off_value=on_off[1])
+                _ = tf.identity(x_, name=_TFOUTPUT)
+                graph = self._run_test_case([_OUTPUT], {_INPUT: x_val})
+                self.assertTrue(len(group_nodes_by_type(graph)["OneHot"]) == 1, "onnx onehot should be used")
 
     @skip_caffe2_backend("issue undefined dim 1")
     def test_flatten0(self):

tf2onnx/tfonnx.py

@@ -1294,8 +1294,9 @@ def onehot_op9(ctx, node, name, args):
     # onnxruntime only supports int64
     output_dtype = ctx.get_dtype(node.input[2])
     if ctx.is_target(TARGET_RS6) and output_dtype not in [onnx_pb.TensorProto.INT64, onnx_pb.TensorProto.INT32]:
-        log.warning("unsupported dtype in onnxruntime, onehot can't be used directly")
+        log.warning("unsupported dtype in onnxruntime, onehot-9 can't be used directly")
         onehot_op(ctx, node, name, args)
+        return
 
     depth = node.input[1]
     depth = ctx.make_node("Unsqueeze", [depth], attr={"axes": [0]}).output[0]
@@ -1307,23 +1308,25 @@ def onehot_op9(ctx, node, name, args):
     off_on_value = ctx.make_node("Concat", [off_value, on_value], attr={"axis": 0}).output[0]
 
     indices = node.input[0]
-    if ctx.get_dtype(indices) != onnx_pb.TensorProto.INT64 and ctx.is_target(TARGET_RS6):
+    if ctx.is_target(TARGET_RS6) and ctx.get_dtype(indices) != onnx_pb.TensorProto.INT64:
         indices = ctx.make_node("Cast", [indices], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
     node.input[0] = indices
 
-    if ctx.get_dtype(depth) != onnx_pb.TensorProto.INT64 and ctx.is_target(TARGET_RS6):
+    if ctx.is_target(TARGET_RS6) and ctx.get_dtype(depth) != onnx_pb.TensorProto.INT64:
         depth = ctx.make_node("Cast", [depth], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
     node.input[1] = depth
 
-    if output_dtype != onnx_pb.TensorProto.INT64 and ctx.is_target(TARGET_RS6):
+    if ctx.is_target(TARGET_RS6) and output_dtype != onnx_pb.TensorProto.INT64:
         off_on_value = ctx.make_node("Cast", [off_on_value], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
     node.input[2] = off_on_value
 
     del node.input[3]
 
-    if output_dtype != onnx_pb.TensorProto.INT64 and ctx.is_target(TARGET_RS6):
-        new_output = utils.make_name("onehot_output")
-        ctx.insert_new_node_on_output("Cast", node.output[0], new_output, to=output_dtype)
+    if ctx.is_target(TARGET_RS6) and output_dtype != onnx_pb.TensorProto.INT64:
+        new_node_name = utils.make_name("onehot_output")
+        new_node = ctx.insert_new_node_on_output("Cast", node.output[0], new_node_name, to=output_dtype)
+        ctx.set_dtype(new_node.output[0], output_dtype)
+        ctx.set_shape(new_node.output[0], ctx.get_shape(node.output[0]))
 
 
 def fused_batchnorm_op7(ctx, node, name, args):