sparse_softmax_cross_entropy_with_logits can be optimized, so don't share function with softmax_cross_entropy_with_logits

sparse_softmax_cross_entropy_with_logits's label must be one-hot, and
optimization can be made according to this attribute

Author: zhijxu-MS

tf2onnx/onnx_opset/nn.py

@@ -706,6 +706,25 @@ def version_7(cls, ctx, node, **kwargs):
         _make_softmax_cross_entropy_with_logits(ctx, labels, logits, node)
 
 
+def _make_sparse_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")
+
+    log_softmax = ctx.make_node(op_type="LogSoftmax", inputs=logit.output)
+    # implement tf.multiply(-1, tf.reduce_sum(tf.multiply(label, log_softmax), axis=1))
+    mul1 = ctx.make_node(op_type="Mul", inputs=[label.output[0], log_softmax.output[0]])
+    reduce_sum = ctx.make_node(op_type="ReduceSum", inputs=[mul1.output[0]], attr={"axes": [-1]})
+    const_negative_one = ctx.make_const(name=utils.make_name("const_negative_one"),
+                                        np_val=np.array(-1).astype(utils.ONNX_TO_NUMPY_DTYPE[logit_dtype]))
+    mul2 = ctx.make_node(op_type="Mul", inputs=[const_negative_one.output[0], reduce_sum.output[0]])
+    shapes = tf_ori_node.output_shapes
+    dtypes = tf_ori_node.output_dtypes
+    ctx.remove_node(tf_ori_node.name)
+    ctx.make_node(op_type="Squeeze", inputs=[mul2.output[0]], attr={"axes": [1]},
+                  outputs=[tf_ori_node.output[0]], shapes=[shapes[0]], dtypes=[dtypes[0]])
+
+
 @tf_op("SparseSoftmaxCrossEntropyWithLogits")
 class SparseSoftmaxCrossEntropyWithLogits:
     @classmethod
@@ -778,4 +797,4 @@ def version_9(cls, ctx, node, **kwargs):
         if logit_dtype != TensorProto.INT64:
             label_node = ctx.make_node("Cast", label_node.output, attr={"to": logit_dtype}, dtypes=[logit_dtype])
 
-        _make_softmax_cross_entropy_with_logits(ctx, label_node, logit_node, node)
+            _make_sparse_softmax_cross_entropy_with_logits(ctx, label_node, logit_node, node)