refactor code

Author: zhijxu-MS

tf2onnx/function/__init__.py

@@ -11,17 +11,17 @@
 from tf2onnx.function.matrixbandpart import matrixbandpart_op
 from tf2onnx.function.range import range_op7
 from tf2onnx.function.select import select_op8
-from tf2onnx.function.sparse_softmax_cross_entropy_with_logits import softmax_cross_entropy_with_logits_op
-from tf2onnx.function.sparse_softmax_cross_entropy_with_logits import sparse_softmax_cross_entropy_with_logits_op
-from tf2onnx.function.sparse_softmax_cross_entropy_with_logits import sparse_softmax_cross_entropy_with_logits_op9
+from tf2onnx.function.softmax_cross_entropy_with_logits import softmax_cross_entropy_with_logits_op7
+from tf2onnx.function.softmax_cross_entropy_with_logits import sparse_softmax_cross_entropy_with_logits_op7
+from tf2onnx.function.softmax_cross_entropy_with_logits import sparse_softmax_cross_entropy_with_logits_op9
 
 __all__ = [
     "gathernd_op",
     "lstm_block_cell_op",
     "matrixbandpart_op",
     "range_op7",
     "select_op8",
-    "softmax_cross_entropy_with_logits_op",
-    "sparse_softmax_cross_entropy_with_logits_op",
+    "softmax_cross_entropy_with_logits_op7",
+    "sparse_softmax_cross_entropy_with_logits_op7",
     "sparse_softmax_cross_entropy_with_logits_op9",
 ]

tf2onnx/function/sparse_softmax_cross_entropy_with_logits.py renamed to tf2onnx/function/softmax_cross_entropy_with_logits.py

@@ -13,7 +13,7 @@
 # pylint: disable=unused-argument,missing-docstring
 
 
-def softmax_cross_entropy_with_logits_computation(ctx, label, logit, tf_ori_node):
+def _make_softmax_cross_entropy_with_logits(ctx, label, logit, tf_ori_node):
     label_dtype = ctx.get_dtype(label.output[0])
     logit_dtype = ctx.get_dtype(logit.output[0])
     utils.make_sure(label_dtype == logit_dtype, "the following logic only works on same dtype of label and logit")
@@ -32,18 +32,18 @@ def softmax_cross_entropy_with_logits_computation(ctx, label, logit, tf_ori_node
                   outputs=[tf_ori_node.output[0]], shapes=[shapes[0]], dtypes=[dtypes[0]])
 
 
-def softmax_cross_entropy_with_logits_op(ctx, node, name, args):
+def softmax_cross_entropy_with_logits_op7(ctx, node, name, args):
     logits = node.inputs[0]
     logit_dtype = ctx.get_dtype(logits.output[0])
     labels = node.inputs[1]
     label_dtype = ctx.get_dtype(labels.output[0])
     if label_dtype != logit_dtype:
         labels = ctx.make_node("Cast", labels.output, attr={"to": logit_dtype}, dtypes=[logit_dtype])
 
-    softmax_cross_entropy_with_logits_computation(ctx, labels, logits, node)
+    _make_softmax_cross_entropy_with_logits(ctx, labels, logits, node)
 
 
-def sparse_softmax_cross_entropy_with_logits_op(ctx, node, name, args):
+def sparse_softmax_cross_entropy_with_logits_op7(ctx, node, name, args):
     # make subgraph to implement one_hot, idea comes from onehot_op
     indices_name = node.input[1]
     indices_shape = ctx.get_shape(indices_name)
@@ -150,4 +150,4 @@ def sparse_softmax_cross_entropy_with_logits_op9(ctx, node, name, args):
     if logit_dtype != TensorProto.INT64:
         label_node = ctx.make_node("Cast", label_node.output, attr={"to": logit_dtype}, dtypes=[logit_dtype])
 
-    softmax_cross_entropy_with_logits_computation(ctx, label_node, logit_node, node)
+    _make_softmax_cross_entropy_with_logits(ctx, label_node, logit_node, node)

tf2onnx/tfonnx.py

@@ -1800,7 +1800,6 @@ def where_op(ctx, node, name, args):
     "ExpandDims": (expanddims_op7, []),
     "OneHot": (onehot_op, []),
     "Reshape": (reshape_op5, []),
-    "SparseSoftmaxCrossEntropyWithLogits": (sparse_softmax_cross_entropy_with_logits_op, [])
 }
 
 _OPSET_6 = {
@@ -1841,7 +1840,8 @@ def where_op(ctx, node, name, args):
     "ResizeNearestNeighbor": (upsample_op7, ["Upsample", "nearest"]),
     "Sin": (direct_op, []),
     "Sub": (broadcast_op7, []),
-    "SoftmaxCrossEntropyWithLogits": (softmax_cross_entropy_with_logits_op, []),
+    "SoftmaxCrossEntropyWithLogits": (softmax_cross_entropy_with_logits_op7, []),
+    "SparseSoftmaxCrossEntropyWithLogits": (sparse_softmax_cross_entropy_with_logits_op7, []),
     "Tan": (direct_op, []),
     "Tile": (tile_op7, []),
     "TruncateDiv": (broadcast_op7, ["Div"]),