Added support for SegmentProd, SegmentMax, SegmentMin

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

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1307,6 +1307,7 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(9, "OneHot")
     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)
@@ -1315,13 +1316,15 @@ def func(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_opset_min_version(9, "OneHot")
+    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)
+            data_val = np.arange(8 * 2 * 3, dtype=np.float32).reshape([8, 2, 3])
+            def func(data, segments):
+                x_ = tf_op(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):

tf2onnx/onnx_opset/reduction.py

@@ -132,7 +132,7 @@ def version_6(cls, ctx, node, **kwargs):
         node.type = "Sum"
 
 
-@tf_op("SegmentSum")
+@tf_op(["SegmentSum", "SegmentProd", "SegmentMax", "SegmentMin"])
 class SegmentSum():
     @classmethod
     def version_9(cls, ctx, node, **kwargs):
@@ -143,20 +143,48 @@ def version_9(cls, ctx, node, **kwargs):
         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))
+        data_is_float = np.dtype(data_np_dtype).kind == 'f'
+        data_is_int = np.dtype(data_np_dtype).kind == 'i'
+        utils.make_sure(data_is_float or data_is_int, "dtype for Segment ops must be float or int")
+
+        if node.type == "SegmentSum":
+            onnx_op = "ReduceSum"
+            identity_value = np.array(0, dtype=data_np_dtype)
+        elif node.type == "SegmentProd":
+            onnx_op = "ReduceProd"
+            identity_value = np.array(1, dtype=data_np_dtype)
+        elif node.type == "SegmentMax":
+            onnx_op = "ReduceMax"
+            if data_is_float:
+                identity_value = np.array('-inf', dtype=data_np_dtype)
+            else:
+                identity_value = np.iinfo(data_np_dtype).min
+        elif node.type == "SegmentMin":
+            onnx_op = "ReduceMin"
+            if data_is_float:
+                identity_value = np.array('inf', dtype=data_np_dtype)
+            else:
+                identity_value = np.iinfo(data_np_dtype).max
+
         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))
+        identity_const = ctx.make_const(utils.make_name("const_identity"), identity_value)
         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})
-        one_hot_unsqueeze = one_hot_node
+        # ORT doesn't support bool for OneHot so we use float32 and cast to bool
+        onehot_values = ctx.make_const(utils.make_name("onehot_values"), np.array([0, 1], dtype=np.float32))
+        one_hot_node = ctx.make_node("OneHot", [segment_inp, num_segments.output[0], onehot_values.output[0]],
+                                     attr={'axis': 0})
+        one_hot_bool = ctx.make_node("Cast", [one_hot_node.output[0]], attr={"to": onnx_pb.TensorProto.BOOL})
+        one_hot_unsqueeze = one_hot_bool
+
         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})
+            one_hot_unsqueeze = ctx.make_node("Unsqueeze", [one_hot_bool.output[0]], attr={'axes': new_dims})
 
-        mul_node = ctx.make_node("Mul", [data_inp, one_hot_unsqueeze.output[0]])
+        mul_node = ctx.make_node("Where", [one_hot_unsqueeze.output[0], data_inp, identity_const.output[0]])
 
         shapes = node.output_shapes
         dtypes = node.output_dtypes
         ctx.remove_node(node.name)
-        sum_node = ctx.make_node("ReduceSum", [mul_node.output[0]], attr={'axes': [1], 'keepdims': 0},
-                                 name=node.name, outputs=node.output, shapes=shapes, dtypes=dtypes)
+        ctx.make_node(onnx_op, [mul_node.output[0]], attr={'axes': [1], 'keepdims': 0},
+                      name=node.name, outputs=node.output, shapes=shapes, dtypes=dtypes)