support Select with Where of opset 9

Author: wayuanho

tests/test_backend.py

@@ -1548,6 +1548,28 @@ def test_reverse_sequence_time_major(self):
 
     @check_opset_min_version(8, "where")
     def test_where(self):
+        x_val = np.array([1, 2, -3, 4, -5, -6, -7, 8, 9, 0], dtype=np.float32)
+        true_result = np.array([111, 222, 333, 444, 555, 666, 777, 888, 999, 1000],
+                               dtype=np.float32)
+        false_result = np.array([-111, -222, -333, -444, -555, -666, -777, -888, -999, -1000],
+                                dtype=np.float32)
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        picks = tf.where(x > -1, true_result, false_result)
+        _ = tf.identity(picks, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+        tf.reset_default_graph()
+        x_val = np.array(1, dtype=np.float32)
+        true_result = np.array(100, dtype=np.float32)
+        false_result = np.array(-111, dtype=np.float32)
+        x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        picks = tf.where(x > -1, true_result, false_result)
+        _ = tf.identity(picks, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
+    @check_opset_min_version(8, "where")
+    @check_target("rs6", "onnxruntime Where type limitation")
+    def test_where_int32(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],
                                dtype=np.int32)
@@ -1560,59 +1582,59 @@ def test_where(self):
 
     @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)
+        x_val = np.array([1, 2, -3, 4, -5, -6, -7, 8, 9, 0], dtype=np.float32)
         true_result = np.array([[111, 111], [222, 222], [333, 333], [444, 444], [555, 555],
                                 [666, 666], [777, 777], [888, 888], [999, 999], [1000, 1000]],
-                               dtype=np.int32)
+                               dtype=np.float32)
         false_result = np.array([[-111, -111], [-222, -222], [-333, -333], [-444, -444],
                                  [-555, -555], [-666, -666], [-777, -777], [-888, -888],
                                  [-999, -999], [-1000, -1000]],
-                                dtype=np.int32)
-        x = tf.placeholder(tf.int32, [None], name=_TFINPUT)
+                                dtype=np.float32)
+        x = tf.placeholder(tf.float32, [None], name=_TFINPUT)
         picks = tf.where(tf.greater_equal(x, 0), true_result, false_result)
         _ = tf.identity(picks, name=_TFOUTPUT)
 
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
     @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)
+        x_val = np.array([[1, 2, -3, 4, -5, -6, -7, 8, 9, 0]], dtype=np.float32)
         true_result = np.array([[111, 222, 333, 444, 555, 666, 777, 888, 999, 1000]],
-                               dtype=np.int32)
+                               dtype=np.float32)
         false_result = np.array([[-111, -222, -333, -444, -555, -666, -777, -888, -999, -1000]],
-                                dtype=np.int32)
-        x = tf.placeholder(tf.int32, [1, 10], name=_TFINPUT)
+                                dtype=np.float32)
+        x = tf.placeholder(tf.float32, [1, 10], name=_TFINPUT)
         picks = tf.where(tf.greater_equal(x, 0), true_result, false_result)
         _ = tf.identity(picks, name=_TFOUTPUT)
 
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
     @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)
+        x_val = np.array([[[1, 2, -3, 4, -5, -6, -7, 8, 9, 0]]], dtype=np.float32)
         true_result = np.array([[[111, 222, 333, 444, 555, 666, 777, 888, 999, 1000]]],
-                               dtype=np.int32)
+                               dtype=np.float32)
         false_result = np.array([[[-111, -222, -333, -444, -555, -666, -777, -888, -999, -1000]]],
-                                dtype=np.int32)
-        x = tf.placeholder(tf.int32, [1, 1, 10], name=_TFINPUT)
+                                dtype=np.float32)
+        x = tf.placeholder(tf.float32, [1, 1, 10], name=_TFINPUT)
         picks = tf.where(tf.greater_equal(x, 0), true_result, false_result)
         _ = tf.identity(picks, name=_TFOUTPUT)
 
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
     @check_opset_min_version(8, "where")
     def test_where_scalar(self):
-        x_val = np.array(6, dtype=np.int32)
+        x_val = np.array(6, dtype=np.float32)
         true_result = np.array([111, 222, 333, 444, 555, 666, 777, 888, 999, 1000],
-                               dtype=np.int32)
+                               dtype=np.float32)
         false_result = np.array([-111, -222, -333, -444, -555, -666, -777, -888, -999, -1000],
-                                dtype=np.int32)
-        x = tf.placeholder(tf.int32, [], name=_TFINPUT)
+                                dtype=np.float32)
+        x = tf.placeholder(tf.float32, [], name=_TFINPUT)
         picks = tf.where(tf.greater_equal(x, 0), true_result, false_result)
         _ = tf.identity(picks, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
-    @check_opset_min_version(9, "where")
+    @check_opset_min_version(9, "NonZero")
     @check_target("rs6", "onnxruntime Transpose type limitation")
     def test_where_with_cond_only(self):
         for np_type, tf_type in [(np.int32, tf.int32), (np.float32, tf.float32)]:

tf2onnx/function/__init__.py

@@ -10,7 +10,7 @@
 from tf2onnx.function.lstm_block_cell import lstm_block_cell_op
 from tf2onnx.function.matrixbandpart import matrixbandpart_op
 from tf2onnx.function.range import range_op7
-from tf2onnx.function.select import select_op8
+from tf2onnx.function.select import select_op8, select_op9
 from tf2onnx.function.softmax_cross_entropy_with_logits import softmax_cross_entropy_with_logits_op7
 from tf2onnx.function.softmax_cross_entropy_with_logits import sparse_softmax_cross_entropy_with_logits_op7
 from tf2onnx.function.softmax_cross_entropy_with_logits import sparse_softmax_cross_entropy_with_logits_op9
@@ -21,6 +21,7 @@
     "matrixbandpart_op",
     "range_op7",
     "select_op8",
+    "select_op9",
     "softmax_cross_entropy_with_logits_op7",
     "sparse_softmax_cross_entropy_with_logits_op7",
     "sparse_softmax_cross_entropy_with_logits_op9",

tf2onnx/function/select.py

@@ -13,6 +13,26 @@
 # pylint: disable=unused-argument,missing-docstring
 
 
+def select_op9(ctx, node, name, args):
+    # T output = Select(bool condition, T x, T y)
+    # T1 output = Where(bool condition, T1 x, T1 y)
+    # NOTE: condition can be 1-dimension in tensorflow, while in onnx,
+    # it should be broadcastable with other two inputs
+    cond_shape = ctx.get_shape(node.input[0])
+    make_sure(cond_shape is not None, "shape of {} is None".format(node.input[0]))
+    input_shape = ctx.get_shape(node.input[1])
+    if input_shape is None:
+        input_shape = ctx.get_shape(node.input[2])
+    make_sure(input_shape is not None, "input shape of {} is None".format(node.name))
+    input_rank = len(input_shape)
+    # if cond shape is 1-dimensional while input has higher rank, need to be reshaped to broadcast
+    if len(cond_shape) == 1 and input_rank > 1:
+        broadcast_shape = [cond_shape[0]] + [1] * (input_rank - 1)
+        shape_const = ctx.make_const(utils.make_name(name), np.array(broadcast_shape, dtype=np.int64))
+        reshape = ctx.make_node("Reshape", [node.input[0], shape_const.output[0]])
+        ctx.replace_input(node, node.input[0], reshape.output[0])
+
+
 def select_op8(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)

tf2onnx/tfonnx.py

@@ -1171,20 +1171,25 @@ def minmax_op(ctx, node, name, args):
         onnx_pb.TensorProto.DOUBLE
     ]
     target_dtype = onnx_pb.TensorProto.FLOAT
+    need_cast = False
     for inp in node.input:
         dtype = ctx.get_dtype(inp)
-        utils.make_sure(dtype, "dtype of {} is None".format(inp))
+        utils.make_sure(dtype is not None, "dtype of {} is None".format(inp))
         if dtype not in supported_dtypes:
             inp_cast = ctx.insert_new_node_on_input(node, "Cast", inp, to=target_dtype)
             ctx.copy_shape(inp, inp_cast.output[0])
             ctx.set_dtype(inp_cast.output[0], target_dtype)
-    origin_dtype = ctx.get_dtype(node.output[0])
-    utils.make_sure(origin_dtype is not None, "dtype of {} is None".format(node.output[0]))
-    ctx.set_dtype(node.output[0], target_dtype)
-    cast_name = utils.make_name(name)
-    cast_node = ctx.insert_new_node_on_output("Cast", node.output[0], name=cast_name, to=origin_dtype)
-    to_replace = [n for n in ctx.get_nodes() if n != cast_node]
-    ctx.replace_all_inputs(to_replace, node.output[0], cast_node.output[0])
+            need_cast = True
+    if need_cast:
+        origin_dtype = ctx.get_dtype(node.output[0])
+        utils.make_sure(origin_dtype is not None, "dtype of {} is None".format(node.output[0]))
+        ctx.set_dtype(node.output[0], target_dtype)
+        cast_name = utils.make_name(name)
+        cast_node = ctx.insert_new_node_on_output("Cast", node.output[0], name=cast_name, to=origin_dtype)
+        ctx.set_dtype(cast_node.output[0], origin_dtype)
+        ctx.copy_shape(node.output[0], cast_node.output[0])
+        to_replace = [n for n in ctx.get_nodes() if n != cast_node]
+        ctx.replace_all_inputs(to_replace, node.output[0], cast_node.output[0])
 
     shapeo = ctx.get_shape(node.output[0])
     needs_broadcast_op = []
@@ -1912,6 +1917,7 @@ def where_op(ctx, node, name, args):
     "ResizeBilinear": (upsample_op9, ["Upsample", "linear"]),
     "ResizeNearestNeighbor": (upsample_op9, ["Upsample", "nearest"]),
     "ReverseSequence": (reverse_op9, []),
+    "Select": (select_op9, ["Where"]),
     "Sign": (sign_op9, []),
     "Sinh": (direct_op, []),
     "SparseSoftmaxCrossEntropyWithLogits": (sparse_softmax_cross_entropy_with_logits_op9, []),
@@ -2209,6 +2215,7 @@ def rewrite_incomplete_type_support(g, ops, impacted_ops):
     """
     ignored_input_index = {
         "Tile": [1],  # Tile's second input can only be int64
+        "Where": [0],  # Where's first input is bool
     }
     new_ops = []
     org_ops = [n for n in ops]
@@ -2264,7 +2271,16 @@ def rewrite_incomplete_type_support_rs5(g, ops):
 
 
 def rewrite_incomplete_type_support_rs6(g, ops):
-    return rewrite_incomplete_type_support(g, ops, ["Div", "IsNaN", "ReduceSum", "Slice", "Split", "Tile", "Transpose"])
+    return rewrite_incomplete_type_support(g, ops, [
+        "Div",
+        "IsNaN",
+        "ReduceSum",
+        "Slice",
+        "Split",
+        "Tile",
+        "Transpose",
+        "Where"
+    ])
 
 
 def rewrite_conv2d_with_pad(g, ops):