enhance tf.where support

tf.where can accept only one input and it will return indices of elem with
true value.

Author: zhijxu-MS

tests/test_backend.py

@@ -1502,6 +1502,16 @@ def test_where_scalar(self):
         _ = tf.identity(picks, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(9, "where")
+    def test_where_with_cond_only(self):
+        for np_type, tf_type in [(np.int32, tf.int32), (np.float32, tf.float32)]:
+            x_val = np.random.randint(0, 2, size=[10, 20, 30]).astype(np_type)
+            x = tf.placeholder(tf_type, shape=[None] * x_val.ndim, name=_TFINPUT)
+            res = tf.where(x)
+            _ = tf.identity(res, name=_TFOUTPUT)
+            self._run_test_case([_OUTPUT], {_INPUT: x_val})
+            tf.reset_default_graph()
+
     @check_opset_min_version(6, "cast")
     def test_shape_int32(self):
         x_val = np.array([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]], dtype=np.float32)

tf2onnx/tfonnx.py

@@ -1713,6 +1713,15 @@ def logical_compare_op(ctx, node, name, args):
     return nodes
 
 
+def where_op(ctx, node, name, args):
+    # T_y output = Where(T_x condition), return indices of elements whose value are True
+    node.type = "NonZero"
+    transpose_node = ctx.insert_new_node_on_output("Transpose", node.output[0], name=utils.make_name("where_op_added"))
+    ctx.set_shape(transpose_node.output[0], ctx.get_shape(node.output[0]))
+    ctx.set_dtype(transpose_node.output[0], ctx.get_dtype(node.output[0]))
+    return [node, transpose_node]
+
+
 # map tensorflow ops to onnx ops. The format below is
 # "TFOP": func_to_map, ["OnnxOp", ...]
 #
@@ -1884,6 +1893,7 @@ def logical_compare_op(ctx, node, name, args):
     "Less": (logical_compare_op, []),
     "ResizeBilinear": (upsample_op9, ["Upsample", "linear"]),
     "ResizeNearestNeighbor": (upsample_op9, ["Upsample", "nearest"]),
+    "Where": (where_op, []),
 }
 
 _OPSETS = [