Added conversion for DenseBincount (#1401)

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

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -4041,6 +4041,20 @@ def func(x):
             return y_
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @skip_tflite("Bug in tflite output shapes")
+    @check_opset_min_version(11, "Unique")
+    @check_tf_min_version("2.3", "needs tf.math.bincount with axis attr")
+    def test_dense_bincount(self):
+        x_val = np.array([[5, 2, 3, 1, 3], [2, 7, 5, 9, 10]], dtype=np.int32)
+        y_val = np.array([[2.0, 1.5, 3.5, 4.5, 5.5], [6.5, 7.5, 8.5, 9.5, 10.5]], dtype=np.float32)
+        for a in [0, -1]:
+            for b in [True, False]:
+                def func(x, y):
+                    x_ = tf.math.bincount(x, axis=a, binary_output=b)
+                    y_ = tf.identity(x_, name=_TFOUTPUT)
+                    return y_
+                self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val})
+
     @check_opset_min_version(11, "ScatterND")
     def test_sparse_to_dense(self):
         i_val = np.array([[0, 0, 0], [0, 0, 2], [0, 1, 3], [1, 2, 2], [1, 2, 3]], dtype=np.int64)

tf2onnx/onnx_opset/reduction.py

@@ -179,6 +179,8 @@ def any_version(cls, opset, ctx, node, **kwargs):
         data_shape = ctx.get_shape(data_inp)
         data_rank = len(data_shape) if data_shape is not None else None
         data_dtype = ctx.get_dtype(data_inp)
+        seg_rank = ctx.get_rank(segment_inp)
+        utils.make_sure(seg_rank == 1, "Segment ops only supported for segments of rank 1, not %s", seg_rank)
         data_np_dtype = utils.map_onnx_to_numpy_type(data_dtype)
         seg_np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(segment_inp))
 

tf2onnx/onnx_opset/tensor.py

@@ -2238,13 +2238,13 @@ def version_11(cls, ctx, node, **kwargs):
                 ctx.copy_shape(new_node.output[2], cast_node.output[0])
 
 
-@tf_op("Bincount")
+@tf_op(["Bincount", "DenseBincount"])
 class Bincount:
     @classmethod
     def any_version(cls, opset, ctx, node, **kwargs):
         # arr, size are int32
         arr_inp, size_inp, weights_inp = node.input
-
+        binary_output = node.get_attr_value("binary_output", False)
         arr_int64 = ctx.make_node("Cast", [arr_inp], attr={'to': TensorProto.INT64}).output[0]
         size_int64 = ctx.make_node("Cast", [size_inp], attr={'to': TensorProto.INT64}).output[0]
 
@@ -2253,22 +2253,55 @@ def any_version(cls, opset, ctx, node, **kwargs):
         weights_is_zero = weights_shape is not None and 0 in weights_shape
         utils.make_sure(weights_is_zero, "Non-empty weights not yet supported for bincount")
 
-        values, _, _, counts = ctx.make_node("Unique", [arr_int64], attr={'sorted': 1}, output_count=4,
-                                             op_name_scope=node.name).output
+        if ctx.get_rank(arr_inp) == 2:
+            zero_const = ctx.make_const(utils.make_name("zero_const"), np.array(0, np.int64)).output[0]
+            one_const = ctx.make_const(utils.make_name("one_const"), np.array(1, np.int64)).output[0]
+            inp_shape = ctx.make_node("Shape", [arr_inp]).output[0]
+            num_rows = GraphBuilder(ctx).make_slice({"data": inp_shape, "starts": [0], "ends": [1], "axes": [0]})
+            num_rows_sq = GraphBuilder(ctx).make_squeeze({"data": num_rows, "axes": [0]})
+            row_idx = ctx.make_node("Range", [zero_const, num_rows_sq, one_const]).output[0]
+            row_idx_unsq = GraphBuilder(ctx).make_unsqueeze({"data": row_idx, "axes": [1]})
+            row_idx_expand = ctx.make_node("Expand", [row_idx_unsq, inp_shape]).output[0]
+            arr_int64_unsq = GraphBuilder(ctx).make_unsqueeze({"data": arr_int64, "axes": [2]})
+            row_idx_expand_unsq = GraphBuilder(ctx).make_unsqueeze({"data": row_idx_expand, "axes": [2]})
+            concat = ctx.make_node("Concat", [row_idx_expand_unsq, arr_int64_unsq], {"axis": 2}).output[0]
+            reshape_const = ctx.make_const(utils.make_name("reshape_const"), np.array([-1, 2], np.int64)).output[0]
+            reshaped = ctx.make_node("Reshape", [concat, reshape_const]).output[0]
+            values, _, _, counts = ctx.make_node("Unique", [reshaped], attr={'sorted': 1, 'axis': 0}, output_count=4,
+                                                 op_name_scope=node.name).output
+            values_to_check_unsq = GraphBuilder(ctx).make_slice(
+                {"data": values, "starts": [1], "ends": [2], "axes": [1]})
+            values_to_check = GraphBuilder(ctx).make_squeeze({"data": values_to_check_unsq, "axes": [1]})
+            size_unsq = GraphBuilder(ctx).make_unsqueeze({'data': size_int64, "axes": [0]})
+            output_shape = ctx.make_node("Concat", [num_rows, size_unsq], attr={"axis": 0}).output[0]
+        else:
+            values, _, _, counts = ctx.make_node("Unique", [arr_int64], attr={'sorted': 1}, output_count=4,
+                                                 op_name_scope=node.name).output
+            values_to_check = values
+            output_shape = GraphBuilder(ctx).make_unsqueeze({'data': size_int64, "axes": [0]})
+
         neg_one_const = ctx.make_const(utils.make_name("neg_one_const"), np.array(-1, np.int64)).output[0]
-        non_neg_val_locs = ctx.make_node("Greater", [values, neg_one_const]).output[0]
-        small_val_locs = ctx.make_node("Less", [values, size_int64]).output[0]
+        non_neg_val_locs = ctx.make_node("Greater", [values_to_check, neg_one_const]).output[0]
+        small_val_locs = ctx.make_node("Less", [values_to_check, size_int64]).output[0]
         valid_val_locs = ctx.make_node("And", [non_neg_val_locs, small_val_locs]).output[0]
 
         valid_values = ctx.make_node("Compress", [values, valid_val_locs], attr={'axis': 0}).output[0]
-        valid_counts = ctx.make_node("Compress", [counts, valid_val_locs], attr={'axis': 0}).output[0]
-
-        output_shape = GraphBuilder(ctx).make_unsqueeze({'data': size_int64, "axes": [0]})
+        if binary_output:
+            counts_shape = ctx.make_node("Shape", [valid_values]).output[0]
+            counts_shape_1d = GraphBuilder(ctx).make_slice(
+                {"data": counts_shape, "starts": [0], "ends": [1], "axes": [0]})
+            ones_tensor = helper.make_tensor("value", TensorProto.INT64, dims=[1], vals=[1])
+            valid_counts = ctx.make_node("ConstantOfShape", [counts_shape_1d], attr={'value': ones_tensor}).output[0]
+        else:
+            valid_counts = ctx.make_node("Compress", [counts, valid_val_locs], attr={'axis': 0}).output[0]
 
-        false_tensor = helper.make_tensor("value", TensorProto.INT64, dims=[1], vals=[0])
-        zeros = ctx.make_node("ConstantOfShape", [output_shape], attr={'value': false_tensor}).output[0]
+        zero_tensor = helper.make_tensor("value", TensorProto.INT64, dims=[1], vals=[0])
+        zeros = ctx.make_node("ConstantOfShape", [output_shape], attr={'value': zero_tensor}).output[0]
 
-        result = ctx.make_node("ScatterElements", [zeros, valid_values, valid_counts], attr={'axis': 0}).output[0]
+        if ctx.get_rank(arr_inp) == 2:
+            result = ctx.make_node("ScatterND", [zeros, valid_values, valid_counts]).output[0]
+        else:
+            result = ctx.make_node("ScatterElements", [zeros, valid_values, valid_counts], attr={'axis': 0}).output[0]
         result_cast = result
         if res_dtype != TensorProto.INT64:
             result_cast = ctx.make_node("Cast", [result], attr={'to': res_dtype}).output[0]