Repository: pytorch/TensorRT

Issue:
Unclear comparisons in conditional operator implementation
<!-- Edit the body of your new issue then click the ✓ "Create Issue" button in the top right of the editor. The first line will be the issue title. Assignees and Labels follow after a blank
line. Leave an empty line before beginning the body of the issue. -->

https://github.com/pytorch/TensorRT/blob/918e9832207d462c2a3aa42f9e7d3ab7aa7415aa/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py#L68
https://github.com/pytorch/TensorRT/blob/918e9832207d462c2a3aa42f9e7d3ab7aa7415aa/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py#L75-L76
https://github.com/pytorch/TensorRT/blob/918e9832207d462c2a3aa42f9e7d3ab7aa7415aa/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py

The above lines are being flagged by mypy as the following:

```sh
py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py:68: error: Non-overlapping equality check (left operand type: "list[Any]", right operand type: "int")  [comparison-overlap]
py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py:76: error: Non-overlapping equality check (left operand type: "list[Any]", right operand type: "int")  [comparison-overlap]
py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py:98: error: Non-overlapping equality check (left operand type: "list[Any]", right operand type: "int")  [comparison-overlap]
```

I cant really figure out what is being checked here but it is likely a bug.


Discussion:

Resolving patch:
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py
--- a/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py
+++ b/py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py
@@ -65,7 +65,7 @@ def where(
         condition_val = condition_layer.get_output(0)
     else:
         assert condition.dtype == trt.bool, "mask dtype is not bool!"
-        if condition_shape != condition_dim:  # TODO: What is this checking?
+        if len(condition_shape) != condition_dim:
             condition_val = expand(
                 network, target, source_ir, f"{name}_expand", condition, output_shape
             )
@@ -73,7 +73,7 @@ def where(
             condition_val = condition
 
     if type(input) != TRTTensor:
-        if x_shape != input_dim:  # TODO: What is this checking?
+        if x_shape != output_shape:
             # special case where 1 element in input
             if len(input.shape) == 0:
                 input = input.unsqueeze(0)
@@ -95,7 +95,7 @@ def where(
         y_val = get_trt_tensor(network, other, f"{name}_y")
     else:
         y_val = other
-        if y_shape != other_dim:  # TODO: What is this checking?
+        if y_shape != output_shape:
             y_val = expand(
                 network, target, source_ir, f"{name}_y_expand", y_val, output_shape
             )


Test patch:
diff --git a/tests/py/dynamo/converters/test_where_aten.py b/tests/py/dynamo/converters/test_where_aten.py
--- a/tests/py/dynamo/converters/test_where_aten.py
+++ b/tests/py/dynamo/converters/test_where_aten.py
@@ -12,6 +12,7 @@ class TestWhereConverter(DispatchTestCase):
             ("2d_broadcast_condition_xshape_yshape", (2, 2), (2, 1)),
             ("3d_condition_xshape_yshape", (2, 2, 1), (2, 2, 1)),
             ("2d_3d_condition_xshape_yshape", (2, 2), (1, 2, 2)),
+            ("3d_2d_condition_xshape_yshape", (1, 2, 2), (2, 2)),
         ]
     )
     def test_(self, _, x_size, y_size):