refactor code

Author: zhijxu-MS

tf2onnx/tfonnx.py

@@ -1290,12 +1290,12 @@ def onehot_op(ctx, node, name, args):
 def onehot_op9(ctx, node, name, args):
     # T output = OneHot(uint8/int32/int64 input, T depth, T on-value, T off-value, @int axis, @dtype)
     # tf requires that dtype is same as on-value's and off-value's dtype
-    # in ONNX, op's schema is (input, depth, [off-value, on-value], @int axis)
+    # in ONNX, op's schema is (input, depth, value, @int axis), meaning of "value" is [off-value, on-value]
     # onnxruntime only supports int64
     output_dtype = ctx.get_dtype(node.input[2])
-    if output_dtype not in [onnx_pb.TensorProto.INT64, onnx_pb.TensorProto.INT32]:
+    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")
-        return onehot_op(ctx, node, name, args)
+        onehot_op(ctx, node, name, args)
 
     depth = node.input[1]
     depth = ctx.make_node("Unsqueeze", [depth], attr={"axes": [0]}).output[0]
@@ -1307,21 +1307,21 @@ 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:
+    if ctx.get_dtype(indices) != onnx_pb.TensorProto.INT64 and ctx.is_target(TARGET_RS6):
         indices = ctx.make_node("Cast", [indices], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
-        node.input[0] = indices
+    node.input[0] = indices
 
-    if ctx.get_dtype(depth) != onnx_pb.TensorProto.INT64:
+    if ctx.get_dtype(depth) != onnx_pb.TensorProto.INT64 and ctx.is_target(TARGET_RS6):
         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:
+    if output_dtype != onnx_pb.TensorProto.INT64 and ctx.is_target(TARGET_RS6):
         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:
+    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)