Merge pull request #351 from nbcsm/revert

Revert "replace conversion logic of "select" with a simpler one"

Author: pengwa

tests/test_backend.py

@@ -1446,7 +1446,8 @@ def test_reverse_sequence_time_major(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(7, "where")
+    # @unittest.skipIf(OPSET < 8, "supported with opset 8 or better")
+    @unittest.skip("FIXME: the newest onnxruntime wheel hasn't been published to PYPI, so Select op is not supported")
     def test_where(self):
         x_val = np.array([1, 2, -3, 4, -5, -6, -7, 8, 9, 0], dtype=np.int32)
         true_result = np.array([111, 222, 333, 444, 555, 666, 777, 888, 999, 1000],
@@ -1458,7 +1459,7 @@ def test_where(self):
         _ = tf.identity(picks, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(7, "where")
+    @check_opset_min_version(8, "where")
     def test_where_with_two_rank_input(self):
         x_val = np.array([1, 2, -3, 4, -5, -6, -7, 8, 9, 0], dtype=np.int32)
         true_result = np.array([[111, 111], [222, 222], [333, 333], [444, 444], [555, 555],
@@ -1474,7 +1475,7 @@ def test_where_with_two_rank_input(self):
 
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(7, "where")
+    @check_opset_min_version(8, "where")
     def test_where_with_two_rank_condition(self):
         x_val = np.array([[1, 2, -3, 4, -5, -6, -7, 8, 9, 0]], dtype=np.int32)
         true_result = np.array([[111, 222, 333, 444, 555, 666, 777, 888, 999, 1000]],
@@ -1487,7 +1488,7 @@ def test_where_with_two_rank_condition(self):
 
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(7, "where")
+    @check_opset_min_version(8, "where")
     def test_where_with_three_rank_condition(self):
         x_val = np.array([[[1, 2, -3, 4, -5, -6, -7, 8, 9, 0]]], dtype=np.int32)
         true_result = np.array([[[111, 222, 333, 444, 555, 666, 777, 888, 999, 1000]]],
@@ -1500,7 +1501,7 @@ def test_where_with_three_rank_condition(self):
 
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(7, "where")
+    @check_opset_min_version(8, "where")
     def test_where_scalar(self):
         x_val = np.array(6, dtype=np.int32)
         true_result = np.array([111, 222, 333, 444, 555, 666, 777, 888, 999, 1000],

tf2onnx/function/__init__.py

@@ -9,7 +9,7 @@
 from .gathernd import gathernd_op
 from .matrixbandpart import matrixbandpart_op
 from .range import range_op7
-from .select import select_op7
+from .select import select_op8
 from .sparse_softmax_cross_entropy_with_logits import sparse_softmax_cross_entropy_with_logits_op
 
-__all__ = ["gathernd_op", "matrixbandpart_op", "range_op7", "select_op7", "sparse_softmax_cross_entropy_with_logits_op"]
+__all__ = ["gathernd_op", "matrixbandpart_op", "range_op7", "select_op8", "sparse_softmax_cross_entropy_with_logits_op"]

tf2onnx/function/select.py

@@ -4,42 +4,236 @@
 """
 tf2onnx.tf2onnx - select op conversion
 """
-
+import numpy as np
+from onnx.onnx_pb import TensorProto
 from tf2onnx import utils
+from tf2onnx.utils import port_name, make_sure
 
 
 # pylint: disable=unused-argument,missing-docstring
 
 
-def select_op7(ctx, node, name, args):
+def select_op8(ctx, node, name, args):
     # T output = Select(bool condition, T x, T y)
-    select_x_dtype = ctx.get_dtype(node.input[1])
-    select_x_shape = ctx.get_shape(node.input[1])
-    cond = node.inputs[0]
-    cond_shape = ctx.get_shape(cond.output[0])
-    utils.make_sure(select_x_shape is not None and cond_shape is not None, "rank of inputs are needed")
-
-    added_nodes = []
-    true_mask = ctx.make_node("Cast", cond.output, attr={"to": select_x_dtype})
-    cond_not = ctx.make_node("Not", cond.output)
-    false_mask = ctx.make_node("Cast", cond_not.output, attr={"to": select_x_dtype})
-    added_nodes.extend([true_mask, cond_not, false_mask])
-    # the broadcasting rule of select is different with common rule.
-    # for example, shape of input_x is (10), shape of input_y is (10, 2) then common broadcasting rule will fail
-    # while in tf "select", input_x will become (10, 1) and repeat at last dimension
-    # so reshape node is inserted here
-    unsqueeze_dim_num = len(select_x_shape) - len(cond_shape)
-    utils.make_sure(unsqueeze_dim_num >= 0, "dim of select_x must not less than cond")
-    if unsqueeze_dim_num != 0:
-        unsqueeze_dim_start = len(select_x_shape)
-        axes = range(unsqueeze_dim_start-1, unsqueeze_dim_start+unsqueeze_dim_num-1)
-        true_mask = ctx.make_node("Unsqueeze", true_mask.output, attr={"axes": axes})
-        false_mask = ctx.make_node("Unsqueeze", false_mask.output, attr={"axes": axes})
-        added_nodes.extend([true_mask, false_mask])
-
-    select_from_true = ctx.make_node("Mul", [true_mask.output[0], node.input[1]])
-    select_from_false = ctx.make_node("Mul", [false_mask.output[0], node.input[2]])
-    res = ctx.make_node("Add", [select_from_true.output[0], select_from_false.output[0]],
-                        name=node.name, outputs=node.output,
-                        shapes=[ctx.get_shape(node.output[0])], dtypes=[ctx.get_dtype(node.output[0])])
-    return [*added_nodes, select_from_true, select_from_false, res]
+    # V v_final_and_scan_outputs = Loop(int64 M, B cond, V v_initial)
+    utils.make_sure(len(node.input) > 1, "Select with only condition is not supported.")
+
+    nodes = []
+    true_data_type = ctx.get_dtype(node.input[1])
+    true_data_shape = ctx.get_shape(node.input[1])
+    make_sure(true_data_type is not None, "select true data dtype cannot be None")
+    make_sure(true_data_shape is not None, "select true data shape cannot be None")
+
+    condition_shape = ctx.get_shape(node.input[0])
+    utils.make_sure(condition_shape is not None, "condition shape is None")
+    rank = len(condition_shape)
+
+    utils.make_sure(rank >= 0, "rank should be >= 0")
+    val_output_id = None
+    if rank > 0:
+        # create nodes getting shape of condition
+        shape_node_output_shape = [rank]
+        shape_node = ctx.make_node("Shape", [node.input[0]], op_name_scope=node.name,
+                                   shapes=[shape_node_output_shape], dtypes=[TensorProto.INT64])
+        nodes.append(shape_node)
+
+        # todo(pengwa), move those leveraging rewrite_incomplete_type_support_onnxruntime after shape inferencing
+        # bug is fixed.
+        # workaround: onnxruntime does not support Split-2, add cases before and after.
+        target_dtype = TensorProto.FLOAT
+        shape_f_node = ctx.make_node("Cast", [shape_node.output[0]], attr={"to": target_dtype},
+                                     shapes=[shape_node_output_shape], dtypes=[target_dtype],
+                                     op_name_scope=node.name)
+        nodes.append(shape_f_node)
+
+        split_attr = [1 for i in range(rank)]
+        output_shapes = [[1] for i in range(rank)]
+        output_dtypes = [target_dtype for i in range(rank)]
+        split_node = ctx.make_node("Split", [shape_f_node.output[0]], output_count=rank,
+                                   attr={"split": split_attr}, shapes=output_shapes,
+                                   dtypes=output_dtypes, op_name_scope=node.name)
+        nodes.append(split_node)
+
+        trip_cnts = []
+        for i in range(rank):
+            output_id = split_node.output[i]
+            output_shape = ctx.get_shape(output_id)
+            target_dtype = TensorProto.INT64
+            shape_i_node = ctx.make_node("Cast", [output_id], attr={"to": target_dtype},
+                                         shapes=[output_shape], dtypes=[target_dtype],
+                                         op_name_scope=node.name)
+            trip_cnts.append(shape_i_node.output[0])
+            nodes.append(shape_i_node)
+        # workaround ends
+
+        onnx_nodes = create_loop_op(ctx, node.input, true_data_type, true_data_shape, trip_cnts, rank)
+        nodes.extend(onnx_nodes)
+        loop_node = onnx_nodes[-1]
+
+        val_output_id = loop_node.output[1]
+    elif rank == 0:
+        if_node, val_output_id = create_if_op(ctx, node.input, true_data_type, true_data_shape)
+        nodes.append(if_node)
+
+    ctx.copy_shape(node.output[0], val_output_id)
+    ctx.set_dtype(node.output[0], true_data_type)
+
+    output_node = ctx.make_node("Identity", [val_output_id], outputs=node.output,
+                                shapes=[ctx.get_shape(val_output_id)], dtypes=[true_data_type])
+    nodes.append(output_node)
+
+    return nodes
+
+
+# gather_input_ids is 1-D tensor, containing 3 elements:
+# 0: condition data to gather on
+# 1: true result to gather on
+# 2: false result to father on
+def create_loop_op(g, gather_input_ids, output_type, output_shape, trip_count_input_ids, rank):
+    nodes = []
+    cond_var_name = utils.make_name("cond_var")
+    nodes.append(g.make_const(cond_var_name, np.array(True, dtype=np.bool)))
+
+    # Loop requires at least a variable, add a useless fake variable.
+    fake_val_name = utils.make_name("fake_var")
+    nodes.append(g.make_const(fake_val_name, np.array(0.0, dtype=np.float32)))
+
+    if rank < 1:
+        raise ValueError("rank is < 1")
+    trip_count_input_id = trip_count_input_ids[-1 * rank]
+
+    loop_inputs = [trip_count_input_id,  # trip count
+                   cond_var_name,  # termination condition
+                   fake_val_name  # initial value of loop-carried dependencies
+                  ]
+    # define an extra scan output
+    loop_node = g.make_node("Loop", loop_inputs, output_count=2, op_name_scope="select_loop",
+                            skip_conversion=False)
+    loop_body = create_loop_body_graph(g, gather_input_ids, output_type, output_shape, trip_count_input_ids,
+                                       rank, loop_node.name)
+    loop_node.set_body_graph_as_attr("body", loop_body)
+    nodes.append(loop_node)
+    return nodes
+
+
+def get_inputs_for_current_iteration(g, input_id, iter_index):
+    nodes = []
+    cond_gather_node = g.make_node("Gather", [input_id, iter_index])
+    nodes.append(cond_gather_node)
+
+    cur_cond_val_scalar_node = g.make_node("Squeeze", [cond_gather_node.output[0]], attr={"axes": [0]})
+    nodes.append(cur_cond_val_scalar_node)
+
+    return nodes, cur_cond_val_scalar_node.output[0]
+
+
+def create_loop_body_graph(parent_g, gather_input_ids, output_data_type, output_shape, trip_count_input_ids,
+                           rank, loop_name):
+    g = parent_g.create_new_graph_with_same_config()
+    iter_name = utils.make_name("i")
+    cond_name = utils.make_name("cond")
+    fake_var_name = utils.make_name("fake_var")
+
+    g.add_graph_input(iter_name, TensorProto.INT64, (1,))  # iteration_num
+    g.add_graph_input(cond_name, TensorProto.BOOL, ())  # condition
+    g.add_graph_input(fake_var_name, TensorProto.FLOAT, ())  # loop-carried dependency
+    nodes = g.get_nodes()
+    # get the i'th value of condition
+    cond_input_id = gather_input_ids[0]
+    new_nodes, cond_input_id_for_current_iter = get_inputs_for_current_iteration(g, cond_input_id, iter_name)
+    nodes.extend(new_nodes)
+
+    # get the i'th value of true values
+    true_input_id = gather_input_ids[1]
+    new_nodes, true_input_id_for_current_iter = get_inputs_for_current_iteration(g, true_input_id, iter_name)
+    nodes.extend(new_nodes)
+
+
+    # get the i'th value of false values
+    false_input_id = gather_input_ids[2]
+    new_nodes, false_input_id_for_current_iter = get_inputs_for_current_iteration(g, false_input_id, iter_name)
+    nodes.extend(new_nodes)
+
+    input_ids_for_current_iter = [cond_input_id_for_current_iter, true_input_id_for_current_iter,
+                                  false_input_id_for_current_iter]
+    output_id = None
+    rank = rank - 1
+    if rank >= 1:
+        nodes_1 = create_loop_op(g, input_ids_for_current_iter, output_data_type, output_shape[1:],
+                                 trip_count_input_ids, rank)
+        loop_1 = nodes_1[-1]
+        output_id = loop_1.output[1]
+        nodes.extend(nodes_1)
+    elif rank == 0:
+        if_node, if_node_output_id = create_if_op(g, input_ids_for_current_iter, output_data_type, output_shape[1:])
+        output_id = if_node_output_id
+        nodes.append(if_node)
+
+    output_identity_name = utils.make_name("loop_output")
+    loop_output_id = utils.port_name(output_identity_name)
+    loop_output_node = g.make_node(
+        'Identity',
+        [output_id],
+        outputs=[loop_output_id],
+        name=output_identity_name
+    )
+    nodes.append(loop_output_node)
+
+    cond_identity_name = utils.make_name("cond_output")
+    cond_output_id = utils.port_name(cond_identity_name)
+    identity_node = g.make_node(
+        'Identity',
+        [cond_name],
+        outputs=[cond_output_id],
+        name=cond_identity_name
+    )
+    nodes.append(identity_node)
+
+    fake_var_identity_name = utils.make_name("fake_var_output")
+    fake_var_output_id = utils.port_name(fake_var_identity_name)
+    identity_node = g.make_node(
+        'Identity',
+        [fake_var_name],
+        outputs=[fake_var_output_id],
+        name=fake_var_identity_name
+    )
+    nodes.append(identity_node)
+
+    g.set_nodes(nodes)
+
+    g.add_graph_output(cond_output_id, TensorProto.BOOL, ())
+    g.add_graph_output(fake_var_output_id, TensorProto.FLOAT, ())
+
+    # use None for all dims, just keep original rank. Because it is observed, dims might be changed in loop.
+    g.add_graph_output(loop_output_id, output_data_type, utils.create_vague_shape_like(output_shape[1:]))
+
+    return g
+
+
+def create_if_op(g, input_ids, output_data_type, output_shape):
+    op_name = utils.make_name("If")
+    true_graph = create_body_graph_for_if_branch(g, output_data_type, output_shape, input_ids[1], op_name)
+    false_graph = create_body_graph_for_if_branch(g, output_data_type, output_shape, input_ids[2], op_name)
+    out_name = port_name(op_name)
+
+    # output a scalar
+    if_node = g.make_node("If", [input_ids[0]], outputs=[out_name], name=op_name, skip_conversion=False)
+    if_node.set_body_graph_as_attr("then_branch", true_graph)
+    if_node.set_body_graph_as_attr("else_branch", false_graph)
+    return if_node, out_name
+
+
+def create_body_graph_for_if_branch(parent_g, data_type, output_shape, chosen_cur_cond_val_out_name, op_name):
+    g = parent_g.create_new_graph_with_same_config()
+    nodes = []
+    name = utils.make_name("Identity")
+    identity_node = g.make_node(
+        'Identity',
+        inputs=[chosen_cur_cond_val_out_name],
+        outputs=['y'],
+        name=name
+    )
+    nodes.append(identity_node)
+    g.set_nodes(nodes)
+    g.add_graph_output("y", data_type, utils.create_vague_shape_like(output_shape))
+    return g

tf2onnx/tfonnx.py

@@ -1834,11 +1834,11 @@ def where_op(ctx, node, name, args):
     "If": (direct_op, []),
     "Loop": (direct_op, []),
     "Scan": (direct_op, []),
-    "Select": (select_op7, []),
 }
 
 _OPSET_8 = {
     "ReverseSequence": (reverse_op8, []),  # make use of scan
+    "Select": (select_op8, []),
 }
 
 _OPSET_9 = {