Improve perf of SegmentSum (#1463)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/backend_test_base.py

@@ -250,6 +250,8 @@ def get_shape(info):
                 return None
             return [d.dim_value if d.HasField('dim_value') else -1 for d in info.type.tensor_type.shape.dim]
         for info in model_shapes.graph.value_info:
+            if info.name == "":
+                continue
             onnx_shape = get_shape(info)
             tf2onnx_shape = graph.get_shape(info.name)
             if onnx_shape is None:

tests/test_backend.py

@@ -1544,7 +1544,7 @@ def func(x, y):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
 
-    @check_opset_min_version(9, "OneHot")
+    @check_opset_min_version(11, "ScatterND")
     def test_segment_sum_data_vector(self):
         segs_val = np.array([0, 0, 0, 1, 2, 2, 3, 3], dtype=np.int32)
         data_val = np.array([5, 1, 7, 2, 3, 4, 1, 3], dtype=np.float32)
@@ -1553,7 +1553,7 @@ def func(data, segments):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: segs_val})
 
-    @check_opset_min_version(11, "Pad")
+    @check_opset_min_version(11, "ScatterND")
     def test_segment_sum_unknown_rank(self):
         segs_val = np.array([0, 0, 0, 1, 2, 2, 3, 3], dtype=np.int32)
         data_val = np.arange(8 * 2 * 3, dtype=np.float32).reshape([8, 2, 3])
@@ -1568,7 +1568,7 @@ def func(data, segments, data_shape, shape_pad):
         self._run_test_case(func, [_OUTPUT],
                             {_INPUT: data_val, _INPUT1: segs_val, _INPUT2: data_shape_val, _INPUT3: shape_pad_val})
 
-    @check_opset_min_version(9, "OneHot")
+    @check_opset_min_version(11, "ScatterND")
     def test_segment_ops_data_tensor(self):
         for tf_op in [tf.math.segment_sum, tf.math.segment_prod, tf.math.segment_min, tf.math.segment_max]:
             segs_val = np.array([0, 0, 0, 1, 2, 2, 3, 3], dtype=np.int32)
@@ -1578,7 +1578,7 @@ def func(data, segments):
                 return tf.identity(x_, name=_TFOUTPUT)
             self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: segs_val})
 
-    @check_opset_min_version(11, "Pad")
+    @check_opset_min_version(11, "ScatterND")
     @skip_tflite("unknown rank")
     def test_segment_mean_unknown_rank(self):
         segs_val = np.array([0, 0, 0, 1, 2, 2, 3, 3], dtype=np.int32)
@@ -1594,7 +1594,7 @@ def func(data, segments, data_shape, shape_pad):
         self._run_test_case(func, [_OUTPUT],
                             {_INPUT: data_val, _INPUT1: segs_val, _INPUT2: data_shape_val, _INPUT3: shape_pad_val})
 
-    @check_opset_min_version(9, "OneHot")
+    @check_opset_min_version(11, "ScatterND")
     def test_sparse_segment_sum(self):
         data_val = np.arange(8 * 2 * 3, dtype=np.float32).reshape([8, 2, 3])
         indices_val = np.array([2, 0, 1, 3, 5, 4, 3, 5, 5], dtype=np.int32)
@@ -1604,7 +1604,7 @@ def func(data, indices, segments):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: indices_val, _INPUT2: segs_val})
 
-    @check_opset_min_version(9, "OneHot")
+    @check_opset_min_version(11, "ScatterND")
     def test_sparse_segment_mean(self):
         data_val = np.arange(8 * 2 * 3, dtype=np.float32).reshape([8, 2, 3])
         indices_val = np.array([2, 0, 1, 3, 5, 4, 3, 5, 5], dtype=np.int32)
@@ -1614,7 +1614,7 @@ def func(data, indices, segments):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: indices_val, _INPUT2: segs_val})
 
-    @check_opset_min_version(9, "OneHot")
+    @check_opset_min_version(11, "ScatterND")
     def test_sparse_segment_sqrtn(self):
         data_val = np.arange(8 * 2 * 3, dtype=np.float32).reshape([8, 2, 3])
         indices_val = np.array([2, 0, 1, 3, 5, 4, 3, 5, 5], dtype=np.int32)
@@ -1624,7 +1624,7 @@ def func(data, indices, segments):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: indices_val, _INPUT2: segs_val})
 
-    @check_opset_min_version(9, "OneHot")
+    @check_opset_min_version(11, "ScatterND")
     def test_sparse_segment_ops_with_num_segments(self):
         for tf_op in [tf.sparse.segment_sum, tf.sparse.segment_mean, tf.sparse.segment_sqrt_n]:
             data_val = np.arange(8 * 2 * 3, dtype=np.float32).reshape([8, 2, 3])
@@ -1635,7 +1635,7 @@ def func(data, indices, segments):
                 return tf.identity(x_, name=_TFOUTPUT)
             self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: indices_val, _INPUT2: segs_val})
 
-    @check_opset_min_version(9, "OneHot")
+    @check_opset_min_version(11, "ScatterND")
     @check_tf_min_version("2.3", "needs tf 2.3")
     def test_unsorted_segment_ops(self):
         tf_ops = [
@@ -1654,7 +1654,7 @@ def func(data, segments):
                 return tf.identity(x_, name=_TFOUTPUT)
             self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: segs_val})
 
-    @check_opset_min_version(9, "OneHot")
+    @check_opset_min_version(11, "ScatterND")
     @check_tf_min_version("2.3", "num_segments can be int64 in tf 2.3")
     def test_segment_op_types(self):
         data_dtypes = [np.int32, np.float32]

tests/test_cond.py

@@ -233,6 +233,7 @@ def case_graph():
 
     @check_tf_min_version("1.8", "shape inference for Reshape op screws up")
     @check_opset_min_version(9, "ConstantOfShape")
+    @allow_missing_shapes("ONNX shape inference still determines if/else shape for unknown reason")
     def test_cond_with_different_output_shape(self):
         input_shape = (10, 5, 20)
         def func(inputs, shape):

tf2onnx/graph_builder.py

@@ -82,7 +82,7 @@ def make_slice(self, kwargs, name=None, shapes=None, dtypes=None, return_node=Fa
             return node
         return node.output[0]
 
-    def make_reduce_sum(self, kwargs, name=None, shapes=None, dtypes=None):
+    def make_reduce_sum(self, kwargs, name=None, shapes=None, dtypes=None, op_name_scope=None):
         """
         ReduceSum changes its schema at opset 13: it treats some axes as dynamic input
         kwargs: key could be ["data", "axes", "keepdims", "noop_with_empty_axes", "outputs"].
@@ -115,7 +115,8 @@ def make_reduce_sum(self, kwargs, name=None, shapes=None, dtypes=None):
         attr = new_attr
 
         return self.graph.make_node(op_type="ReduceSum", inputs=inputs, attr=attr, name=name,
-                                    outputs=outputs, shapes=shapes, dtypes=dtypes).output[0]
+                                    outputs=outputs, shapes=shapes, dtypes=dtypes,
+                                    op_name_scope=op_name_scope).output[0]
 
     def make_squeeze(self, kwargs, name=None, shapes=None, dtypes=None, return_node=False, op_name_scope=None):
         """

tf2onnx/onnx_opset/controlflow.py

@@ -217,6 +217,7 @@ def version_9(cls, ctx, node, **kwargs):
                                                        node.output[0], name=utils.make_name("where_op_added"))
         ctx.copy_shape(node.output[0], transpose_node.output[0])
         ctx.copy_dtype(node.output[0], transpose_node.output[0])
+        ctx.update_node_shape_dtype(node, override=True)
 
 
 @tf_op(["StatelessIf"])