replace conversion logic of "select" with a simpler one

Author: zhijxu-MS

tests/test_backend.py

@@ -1446,8 +1446,7 @@ def test_reverse_sequence_time_major(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    # @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")
+    @check_opset_min_version(7, "where")
     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],
@@ -1459,7 +1458,7 @@ def test_where(self):
         _ = tf.identity(picks, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(8, "where")
+    @check_opset_min_version(7, "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],
@@ -1475,7 +1474,7 @@ def test_where_with_two_rank_input(self):
 
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(8, "where")
+    @check_opset_min_version(7, "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]],
@@ -1488,7 +1487,7 @@ def test_where_with_two_rank_condition(self):
 
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(8, "where")
+    @check_opset_min_version(7, "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]]],
@@ -1501,7 +1500,7 @@ def test_where_with_three_rank_condition(self):
 
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(8, "where")
+    @check_opset_min_version(7, "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_op8
+from .select import select_op7
 from .sparse_softmax_cross_entropy_with_logits import sparse_softmax_cross_entropy_with_logits_op
 
-__all__ = ["gathernd_op", "matrixbandpart_op", "range_op7", "select_op8", "sparse_softmax_cross_entropy_with_logits_op"]
+__all__ = ["gathernd_op", "matrixbandpart_op", "range_op7", "select_op7", "sparse_softmax_cross_entropy_with_logits_op"]

tf2onnx/function/select.py

@@ -4,236 +4,42 @@
 """
 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_op8(ctx, node, name, args):
+def select_op7(ctx, node, name, args):
     # T output = Select(bool condition, T x, T y)
-    # 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
+    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]

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 = {