Repository: apache/tvm

Issue:
[Bug] [Frontend][Tensorflow] tf.where with broadcast condition fails to import due to Incompatible broadcast type
The test case below fails to import in tvm:

```
def test_forward_where_with_broadcast_cond():
    t1 = np.array([1.0, 2.0, 3.0, 4.0, 5.0]).astype("float32")
    t2 = np.array([2.0, 4.0, 1.0, 3.0, 5.0]).astype("float32")
    x = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]).astype("float32")
    y = np.array([[10.0, 9.0], [8.0, 7.0], [6.0, 5.0], [4.0, 3.0], [2.0, 1.0]]).astype("float32")

    with tf.Graph().as_default():
        in1 = tf.placeholder(shape=(5), dtype = "float32", name="in1")
        in2 = tf.placeholder(shape=(5), dtype = "float32", name="in2")
        condition = math_ops.less(in1, in2, name="less")
        lhs = tf.placeholder(shape=(5,2), dtype = "float32", name="x")
        rhs = tf.placeholder(shape=(5,2), dtype = "float32", name="y")
        out = tf.where(condition, lhs, rhs)
        compare_tf_with_tvm([t1, t2, x, y], ["in1:0", "in2:0", "x:0", "y:0"], out.name)
```

### Expected behavior
Should be identical to tensorflow output:
```
[array([[1., 2.],
       [3., 4.],
       [6., 5.],
       [4., 3.],
       [2., 1.]], dtype=float32)]
```

### Actual behavior
Failed with this error:

`Incompatible broadcast type TensorType([5], bool) and TensorType([5, 2], float32)`

### Environment
```
Linux
LSB Version:    :core-4.1-amd64:core-4.1-ia32:core-4.1-noarch:cxx-4.1-amd64:cxx-4.1-ia32:cxx-4.1-noarch:desktop-4.1-amd64:desktop-4.1-ia32:desktop-4.1-noarch:languages-4.1-amd64:languages-4.1-noarch:printing-4.1-amd64:printing-4.1-noarch
Distributor ID: RedHatEnterpriseWorkstation
Description:    Red Hat Enterprise Linux Workstation release 7.9 (Maipo)
Release:        7.9
Codename:       Maipo
Name: apache-tvm
Version: 0.10.0
Home-page: https://tlcpack.ai
Author: Apache TVM
Author-email: None
License: Apache
Name: tensorflow
Version: 2.5.0
Summary: TensorFlow is an open source machine learning framework for everyone.
Home-page: https://www.tensorflow.org/
Author: Google Inc.
Author-email: packages@tensorflow.org
License: Apache 2.0
```

### Steps to reproduce
Add above testcase to tests/python/frontend/tensorflow/test_forward.py and run
`python -m pytest tests/python/frontend/tensorflow/test_forward.py -k test_forward_where_with_broadcast_cond`

### Triage

* needs-triage
* frontend:tensorflow 


Discussion:
This test fails with numpy and tensorflow 2.x but works fine with tensorflow 1.x:
```
import numpy as np
import tensorflow as tf

try:
    tf_compat_v1 = tf.compat.v1
except:
    tf_compat_v1 = tf

t1 = np.arange(5, dtype="float32")
t2 = np.arange(5, dtype="float32")[::-1]
x = np.arange(10, dtype="float32").reshape(5,2)
y = np.arange(10, dtype="float32")[::-1].reshape(5,2)
condition = np.less(t1, t2)
print("tf1.x output:")
print(tf_compat_v1.where(condition, x, y))
print("tf2.x output:")
print(tf.where(condition, x, y))
```



Resolving patch:
diff --git a/python/tvm/relay/frontend/tensorflow_ops.py b/python/tvm/relay/frontend/tensorflow_ops.py
--- a/python/tvm/relay/frontend/tensorflow_ops.py
+++ b/python/tvm/relay/frontend/tensorflow_ops.py
@@ -2380,6 +2380,24 @@ def _impl(inputs, attr, params, mod):
 
 
 def _where():
+    def _impl(inputs, attr, params, mod):
+        if len(inputs) == 1:
+            return AttrCvt(op_name="argwhere")(inputs, attr)
+        cond_shape = _infer_shape(inputs[0], mod)
+        x_shape = _infer_shape(inputs[1], mod)
+        # Due to difference in broadcast behavior between Select and SelectV2,
+        # we adjust condition dimension with expand_dim and then broadcast.
+        if len(cond_shape) == 1 and cond_shape[0] == x_shape[0]:
+            for _ in range(len(x_shape) - 1):
+                inputs[0] = _op.expand_dims(inputs[0], axis=-1)
+            broadcast_cond = _op.broadcast_to(inputs[0], x_shape)
+            inputs[0] = _op.cast(broadcast_cond, "bool")
+        return AttrCvt(op_name="where")(inputs, attr)
+
+    return _impl
+
+
+def _where_v2():
     def _impl(inputs, attr, params, mod):
         if len(inputs) == 1:
             return AttrCvt(op_name="argwhere")(inputs, attr)
@@ -3088,7 +3106,7 @@ def _impl(inputs, attr, params, mod):
     "Round": AttrCvt("round"),
     "Rsqrt": _rsqrt(),
     "Select": _where(),
-    "SelectV2": _where(),
+    "SelectV2": _where_v2(),
     "Selu": _selu(),
     "Shape": _shape(),
     "Sigmoid": AttrCvt("sigmoid"),
@@ -3142,6 +3160,6 @@ def _impl(inputs, attr, params, mod):
     "UniqueWithCounts": _unique(True),
     "Unpack": _unpack(),
     "UnravelIndex": _unravel_index(),
-    "Where": _where(),
+    "Where": _where_v2(),
     "ZerosLike": AttrCvt("zeros_like"),
 }


Test patch:
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -1234,6 +1234,29 @@ def test_forward_argwhere():
     _test_forward_where((5, 5, 5, 5, 5))
 
 
+def _test_forward_where_with_broadcast(in_shape, cond_shape):
+    choice_list = list(np.arange(10).astype("float32"))
+    t1 = np.random.choice(choice_list, size=cond_shape)
+    t2 = np.random.choice(choice_list, size=cond_shape)
+    x = np.random.choice(choice_list, size=in_shape)
+    y = np.random.choice(choice_list, size=in_shape)
+
+    with tf.Graph().as_default():
+        in1 = tf.placeholder(shape=cond_shape, dtype="float32", name="in1")
+        in2 = tf.placeholder(shape=cond_shape, dtype="float32", name="in2")
+        condition = math_ops.less(in1, in2, name="less")
+        lhs = tf.placeholder(shape=in_shape, dtype="float32", name="x")
+        rhs = tf.placeholder(shape=in_shape, dtype="float32", name="y")
+        out = tf.where(condition, lhs, rhs)
+        compare_tf_with_tvm([t1, t2, x, y], ["in1:0", "in2:0", "x:0", "y:0"], out.name)
+
+
+def test_forward_where_with_broadcast():
+    _test_forward_where_with_broadcast((5, 2), (5,))
+    _test_forward_where_with_broadcast((5, 7), (5,))
+    _test_forward_where_with_broadcast((3, 2, 5), (3,))
+
+
 #######################################################################
 # SpaceToBatchND
 # --------------