support cumprod op (#2133)

* support cumprod op

---------

Signed-off-by: Francesco Salvetti <[email protected]>
Co-authored-by: Jay Zhang <[email protected]>

Author: f-salvetti

tests/test_backend.py

@@ -4718,6 +4718,53 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(10, "Slice")
+    def test_cumprod(self):
+        x_val = np.array([2.0, 3.0, 4.0, 5.0], dtype=np.float32).reshape((2, 2))
+        def func(x):
+            x_ = tf.math.cumprod(x, axis=0)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+    @check_opset_min_version(10, "Slice")
+    def test_cumprod_axis1(self):
+        x_val = np.array([1., 2., 3., 4.,
+                          5., 6., 7., 8.,
+                          9., 10., 11., 12.,
+                          13., 14., 15., 16.,
+                          17., 18., 19., 20.,
+                          21., 22., 23., 24.], dtype=np.float32).reshape((2, 3, 4))
+        def func(x):
+            x_ = tf.math.cumprod(x, axis=1)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+    @check_opset_min_version(10, "Slice")
+    def test_cumprod_axis1_reverse(self):
+        x_val = np.array([1., 2., 3., 4.,
+                          5., 6., 7., 8.,
+                          9., 10., 11., 12.,
+                          13., 14., 15., 16.,
+                          17., 18., 19., 20.,
+                          21., 22., 23., 24.], dtype=np.float32).reshape((2, 3, 4))
+        def func(x):
+            x_ = tf.math.cumprod(x, axis=1, reverse=True)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
+    @check_opset_min_version(10, "Slice")
+    def test_cumprod_axis1_reverse_exclusive(self):
+        x_val = np.array([1., 2., 3., 4.,
+                          5., 6., 7., 8.,
+                          9., 10., 11., 12.,
+                          13., 14., 15., 16.,
+                          17., 18., 19., 20.,
+                          21., 22., 23., 24.], dtype=np.float32).reshape((2, 3, 4))
+        def func(x):
+            x_ = tf.math.cumprod(x, axis=1, reverse=True, exclusive=True)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
     @check_opset_min_version(11, "Round")
     def test_round(self):
         x_val = np.array([-0.7, -0.5, -0.0, 0.0, +0.0, 0.3, 0.5, 0.7, float('nan')], dtype=np.float32)

tf2onnx/onnx_opset/math.py

@@ -12,6 +12,8 @@
 from tf2onnx import constants, utils
 from tf2onnx.handler import tf_op
 from tf2onnx.onnx_opset import common
+from tf2onnx.graph_builder import GraphBuilder
+
 
 logger = logging.getLogger(__name__)
 
@@ -555,6 +557,207 @@ def version_11(cls, ctx, node, **kwargs):
         pass
 
 
+@tf_op("Cumprod")
+class CumProd:
+    @classmethod
+    def version_10(cls, ctx, node, **kwargs):
+        # opset 10 required for Slice to support starts/ends/axes/steps as inputs
+        axis_node = node.inputs[1]
+        is_axis_const = axis_node.is_const()
+        if is_axis_const: # we can compute axis value right now
+            axis = axis_node.get_tensor_value()
+            axis_node = ctx.make_const(utils.make_name("axis"), np.array([axis], dtype=np.int64))
+        else:
+            axis_node = ctx.make_node("Cast", inputs=[axis_node.output[0]], attr={"to": onnx_pb.TensorProto.INT64},
+                                      op_name_scope=node.name, outputs=[utils.make_name("axis")])
+            axis_node = GraphBuilder(ctx).make_unsqueeze({'data': axis_node.output[0], 'axes': [0]}, return_node=True)
+            axis = axis_node.output[0]
+
+        input_rank = len(ctx.get_shape(node.input[0]))
+        cond_true_node = ctx.make_const(utils.make_name("cond_in"), np.ones((), dtype=bool))
+        input_shape_node = ctx.make_node("Shape", inputs=[node.input[0]], op_name_scope=node.name,
+                                         outputs=[utils.make_name("input_shape")])
+        axis_length_node = ctx.make_node("Gather", inputs=[input_shape_node.output[0], node.input[1]],
+                                         op_name_scope=node.name, outputs=[utils.make_name("axis_length")])
+        one_node = ctx.make_const(utils.make_name("one"), np.array([1], "int64"))
+        axis_length_plus_one_node = ctx.make_node("Add", inputs=[axis_length_node.output[0], one_node.output[0]],
+                                                  op_name_scope=node.name,
+                                                  outputs=[utils.make_name("axis_length_plus_one")])
+        num_iter_node = ctx.make_node("Sub", inputs=[axis_length_node.output[0], one_node.output[0]],
+                                      op_name_scope=node.name, outputs=[utils.make_name("num_iter")])
+
+        if node.get_attr_value("exclusive"): # one iter less, crop the input, then pad the output
+            num_iter_node = ctx.make_node("Sub", inputs=[num_iter_node.output[0], one_node.output[0]],
+                                          op_name_scope=node.name, outputs=[utils.make_name("num_iter")])
+            zero_node = ctx.make_const(utils.make_name("zero"), np.array([0], "int64"))
+            if node.get_attr_value("reverse"):
+                pad_axis = [0, 1]
+                start_slice = one_node.output[0]
+                end_slice = axis_length_plus_one_node.output[0]
+            else:
+                minus_one_node = ctx.make_const(utils.make_name("minus_one"), np.array([-1], "int64"))
+                pad_axis = [1, 0]
+                start_slice = zero_node.output[0]
+                end_slice = minus_one_node.output[0]
+            pads_node = cls.get_pads_node(ctx, pad_axis, axis, input_rank, node.name)
+            slice_shape = [-1] * len(ctx.get_shape(node.input[0]))
+            inputs_node = ctx.make_node("Slice", inputs=[node.input[0], start_slice, end_slice, axis_node.output[0]],
+                                        op_name_scope=node.name, outputs=[utils.make_name("slice")],
+                                        shapes=[slice_shape], dtypes=[ctx.get_dtype(node.input[0])])
+            inputs = inputs_node.output[0]
+        else:
+            inputs = node.input[0]
+
+        loop_graph = cls.make_loop_graph(ctx, node, inputs, input_rank, axis)
+        loop_graph.parent_graph = ctx
+
+        loop_inputs = [num_iter_node.output[0], cond_true_node.output[0], inputs,
+                       axis_length_plus_one_node.output[0], inputs]
+        loop_outputs = [utils.make_name("loop_inputs_out"), utils.make_name("loop_axis_length_plus_one_out"),
+                        utils.make_name("loop_accumulator_out")]
+        if not is_axis_const: # axis is a tensor, we neeed to feed it to the loop graph
+            loop_inputs.append(axis)
+            loop_outputs.append(utils.make_name("loop_axis_out"))
+        loop_outputs_shapes = [loop_graph.get_shape(o) for o in loop_graph.outputs[1:]]
+        loop_outputs_dtypes = [loop_graph.get_dtype(o) for o in loop_graph.outputs[1:]]
+
+        loop_node = ctx.make_node("Loop", inputs=loop_inputs, branches={"body": loop_graph}, outputs=loop_outputs,
+                                  shapes=loop_outputs_shapes, dtypes=loop_outputs_dtypes, op_name_scope=node.name)
+
+        if node.get_attr_value("exclusive"): # pad the output
+            if ctx.get_dtype(loop_node.output[2]) != ctx.get_dtype(one_node.output[0]):
+                pad_const_node = ctx.make_node("Cast", inputs=[one_node.output[0]],
+                                               attr={"to": ctx.get_dtype(loop_node.output[2])},
+                                               op_name_scope=node.name, outputs=[utils.make_name("pad_const")])
+            else:
+                pad_const_node = one_node
+            output_node = ctx.make_node("Pad", op_name_scope=node.name, outputs=[utils.make_name("cumprod_out")],
+                                        inputs=[loop_node.output[2], pads_node.output[0], pad_const_node.output[0]])
+            output = output_node.output[0]
+        else:
+            output = loop_node.output[2]
+        output_node = ctx.make_node("Identity", inputs=[output], outputs=[utils.make_name("cumprod_out")],
+                                    shapes=[ctx.get_shape(node.input[0])], dtypes=[ctx.get_dtype(node.input[0])])
+        ctx.insert_node_on_output(output_node, node.output[0])
+        ctx.remove_node(node.name)
+
+    @classmethod
+    def make_loop_graph(cls, ctx, node, inputs_tensor, input_rank, axis):
+        inputs_tensor_shape = ctx.get_shape(inputs_tensor)
+        inputs_tensor_dtype = ctx.get_dtype(inputs_tensor)
+
+        graph = ctx.create_new_graph_with_same_config()
+        graph.add_graph_input(utils.make_name("iteration_num"), onnx_pb.TensorProto.INT64, [])
+        graph.add_graph_input(utils.make_name("condition_in"), onnx_pb.TensorProto.BOOL, [])
+        graph.add_graph_input(utils.make_name("inputs"), inputs_tensor_dtype, inputs_tensor_shape)
+        graph.add_graph_input(utils.make_name("axis_length_plus_one"), onnx_pb.TensorProto.INT64, [1])
+        graph.add_graph_input(utils.make_name("accumulator"), inputs_tensor_dtype, inputs_tensor_shape)
+        if not isinstance(axis, int): # axis is a tensor, we need to feed it to the loop graph
+            graph.add_graph_input(utils.make_name("axis"), onnx_pb.TensorProto.INT64, [1])
+            axis = graph.input_names[-1]
+            axis_node = graph.get_node_by_output(axis)
+        else:
+            axis_node = graph.make_const(utils.make_name("axis"), np.array([axis], "int64"))
+
+        # main loop graph
+        loop_name = node.name + "/loop"
+        iter_num = GraphBuilder(graph).make_unsqueeze({'data': graph.input_names[0], 'axes': [0]})
+        one_node = graph.make_const(utils.make_name("one"), np.array(1, "int64"))
+        zero_node = graph.make_const(utils.make_name("zero"), np.array([0], "int64"))
+
+        add_node = graph.make_node("Add", inputs=[iter_num, one_node.output[0]],
+                                   outputs=[utils.make_name("add")], op_name_scope=loop_name)
+
+        if node.get_attr_value("reverse"):
+            pad_axis = [zero_node.output[0], add_node.output[0]]
+            start_slice = add_node.output[0]
+            end_slice = graph.input_names[3]
+        else:
+            neg_node = graph.make_node("Neg", inputs=[add_node.output[0]],
+                                       outputs=[utils.make_name("neg")], op_name_scope=loop_name)
+            pad_axis = [add_node.output[0], zero_node.output[0]]
+            start_slice = zero_node.output[0]
+            end_slice = neg_node.output[0]
+
+        pads_node = cls.get_pads_node(graph, pad_axis, axis, input_rank, is_pad_axis_const=False, base_name=loop_name)
+        slice_node = graph.make_node("Slice", op_name_scope=loop_name, outputs=[utils.make_name("slice")],
+                                     inputs=[graph.input_names[2], start_slice, end_slice, axis_node.output[0]])
+        if graph.get_dtype(slice_node.output[0]) != graph.get_dtype(one_node.output[0]):
+            pad_const_node = graph.make_node("Cast", inputs=[one_node.output[0]],
+                                             attr={"to": graph.get_dtype(slice_node.output[0])},
+                                             op_name_scope=loop_name, outputs=[utils.make_name("pad_const")])
+        else:
+            pad_const_node = one_node
+        pad_node = graph.make_node("Pad", inputs=[slice_node.output[0], pads_node.output[0], pad_const_node.output[0]],
+                                   op_name_scope=loop_name, outputs=[utils.make_name("pad")])
+        mul_node = graph.make_node("Mul", inputs=[graph.input_names[4], pad_node.output[0]],
+                                   op_name_scope=loop_name, outputs=[utils.make_name("mul")],
+                                   shapes=[inputs_tensor_shape], dtypes=[inputs_tensor_dtype])
+
+        # manage loop outputs
+        output_cond_node = graph.make_node("Identity", inputs=[graph.input_names[1]], op_name_scope=loop_name,
+                                           outputs=[utils.make_name("condition_out")])
+        output_inp_node = graph.make_node("Identity", inputs=[graph.input_names[2]], op_name_scope=loop_name,
+                                          outputs=[utils.make_name("inputs_out")])
+        output_axis_length_plus_one_node = graph.make_node("Identity", inputs=[graph.input_names[3]],
+                                                           op_name_scope=loop_name,
+                                                           outputs=[utils.make_name("axis_length_plus_one_out")])
+        output_acc_node = graph.make_node("Identity", inputs=[mul_node.output[0]], op_name_scope=loop_name,
+                                          outputs=[utils.make_name("accumulator_out")])
+
+        graph.add_graph_output(output_cond_node.output[0])                  # 1 condition output
+        graph.add_graph_output(output_inp_node.output[0])                   # N loop carried dependencies outputs
+        graph.add_graph_output(output_axis_length_plus_one_node.output[0])  # N loop carried dependencies outputs
+        graph.add_graph_output(output_acc_node.output[0])                   # N loop carried dependencies outputs
+
+        if not isinstance(axis, int): # axis is a tensor, we need to feed it to the loop graph
+            output_axis_node = graph.make_node("Identity", inputs=[axis], op_name_scope=loop_name,
+                                               outputs=[utils.make_name("axis_out")])
+            graph.add_graph_output(output_axis_node.output[0])              # N loop carried dependencies outputs
+        return graph
+
+    @classmethod
+    def get_pads_node(cls, graph, pad_axis, axis, rank, is_pad_axis_const=True, base_name=""):
+        if isinstance(axis, int): # axis, is a const, we directly compute padding values
+            pre_pad = np.zeros(axis, "int64")
+            post_pad = np.zeros(rank - axis - 1, "int64")
+            if is_pad_axis_const: # pylint: disable=R1705
+                pads = np.concatenate([pre_pad, pad_axis[0:1], post_pad,
+                                       pre_pad, pad_axis[1:2], post_pad])
+                pads_node = graph.make_const(utils.make_name("pads"), pads)
+                return pads_node
+            else:
+                pre_pad_node = graph.make_const(utils.make_name("pre_pad"), pre_pad)
+                post_pad_node = graph.make_const(utils.make_name("post_pad"), post_pad)
+
+        else: # axis is a tensor, we need to compute padding values at runtime
+            if is_pad_axis_const:
+                pad_axis = [graph.make_const(utils.make_name("pad"),
+                                             np.array([pad], "int64")).output[0] for pad in pad_axis]
+
+            rank_tensor = graph.make_const(utils.make_name("rank"), np.array([rank], "int64")).output[0]
+            zero_node = graph.make_const(utils.make_name("zero"), np.array([0], "int64"))
+            one_node = graph.make_const(utils.make_name("zero"), np.array([1], "int64"))
+
+            post_repeat_node = graph.make_node("Sub", inputs=[rank_tensor, axis],
+                                               outputs=[utils.make_name("post_repeat")], op_name_scope=base_name)
+            post_repeat_node = graph.make_node("Sub", inputs=[post_repeat_node.output[0], one_node.output[0]],
+                                               outputs=[utils.make_name("post_repeat")], op_name_scope=base_name)
+
+            pre_pad_node = graph.make_node("Tile", inputs=[zero_node.output[0], axis], op_name_scope=base_name,
+                                           attr={"axis": 0}, outputs=[utils.make_name("pre_pad")])
+
+            post_pad_node = graph.make_node("Tile", inputs=[zero_node.output[0], post_repeat_node.output[0]],
+                                            attr={"axis": 0}, outputs=[utils.make_name("post_pad")],
+                                            op_name_scope=base_name)
+
+        pads_node = graph.make_node("Concat", attr={"axis": 0}, outputs=[utils.make_name("pads")],
+                                    op_name_scope=base_name,
+                                    inputs=[pre_pad_node.output[0], pad_axis[0], post_pad_node.output[0],
+                                            pre_pad_node.output[0], pad_axis[1], post_pad_node.output[0]])
+        return pads_node
+
+
 @tf_op("Round")
 class Round:
     @classmethod

tf2onnx/version.py

@@ -2,4 +2,4 @@
 
 
 version = '1.13.0'
-git_version = 'ddca3a5eb2d912f20fe7e0568dd1a3013aee9fa3'
+git_version = '82ae7c4659ab9ece121a6414ee037d68b6b2d907'