Made SegmentSum support higher-dimensional data

Signed-off-by: Tom Wildenhain <tomwi@microsoft.com>

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1307,14 +1307,22 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
-    def test_segment_sum(self):
+    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)
         def func(data, segments):
             x_ = tf.math.segment_sum(data, segments)
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: segs_val})
 
+    def test_segment_sum_data_tensor(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])
+        def func(data, segments):
+            x_ = tf.math.segment_sum(data, segments)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: segs_val})
+
     @check_onnxruntime_incompatibility("Sqrt")
     def test_sqrt(self):
         x_val = np.array([4.0, 16.0, 4.0, 1.6], dtype=np.float32).reshape((2, 2))

tf2onnx/onnx_opset/reduction.py

@@ -140,14 +140,20 @@ def version_9(cls, ctx, node, **kwargs):
         segment_inp = node.input[1]
         data_shape = ctx.get_shape(data_inp)
         utils.make_sure(data_shape is not None, "Segment ops require input rank to be known")
+        data_rank = len(data_shape)
         data_np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(data_inp))
         seg_np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(segment_inp))
         max_segment = ctx.make_node("ReduceMax", [segment_inp], attr={'axes': [0], 'keepdims': 0})
         one_const = ctx.make_const(utils.make_name("const_one"), np.array(1, dtype=seg_np_dtype))
         num_segments = ctx.make_node("Add", [max_segment.output[0], one_const.output[0]])
         onehot_values = ctx.make_const(utils.make_name("onehot_values"), np.array([0, 1], dtype=data_np_dtype))
         one_hot_node = ctx.make_node("OneHot", [segment_inp, num_segments.output[0], onehot_values.output[0]], attr={'axis': 0})
-        mul_node = ctx.make_node("Mul", [data_inp, one_hot_node.output[0]])
+        one_hot_unsqueeze = one_hot_node
+        if data_rank > 1:
+            new_dims = list(range(2, 2 + data_rank - 1))
+            one_hot_unsqueeze = ctx.make_node("Unsqueeze", [one_hot_node.output[0]], attr={'axes': new_dims})
+
+        mul_node = ctx.make_node("Mul", [data_inp, one_hot_unsqueeze.output[0]])
 
         shapes = node.output_shapes
         dtypes = node.output_dtypes