Repository: pytorch/TensorRT

Issue:
🐛 [Bug] Issue in `broadcastable` converter utility
##  Bug Description

Encountered `IndexError` in `broadcastable` utility:

```python
File "~/TensorRT/py/torch_tensorrt/dynamo/conversion/converter_utils.py", line 118, in broadcastable
     greater_tensor[i] != lesser_tensor[j]

IndexError: tuple index out of range

While executing %where_5 : [num_users=1] = call_function[target=torch.ops.aten.where.self](...)
```

## To Reproduce

Steps to reproduce the behavior:

Invoke `aten.where` converter with a 0D tensor.

## Expected behavior

Utility should not error.

## Environment

 - Torch-TensorRT Version (e.g. 1.0.0): 8c62fca7649ac2e8e74ab7d9d66a20a164e2af51
 - PyTorch Version (e.g. 1.0): `2.1.0.dev20230803+cu121`

## Additional context

Blocking compilation of #1697 with torch compile path.



Discussion:

Resolving patch:
diff --git a/py/torch_tensorrt/dynamo/conversion/converter_utils.py b/py/torch_tensorrt/dynamo/conversion/converter_utils.py
--- a/py/torch_tensorrt/dynamo/conversion/converter_utils.py
+++ b/py/torch_tensorrt/dynamo/conversion/converter_utils.py
@@ -137,25 +137,23 @@ def broadcastable(
     "Check if two tensors are broadcastable according to torch rules"
     a_shape = tuple(a.shape)
     b_shape = tuple(b.shape)
+
     # check from the trailing
     diff = len(a_shape) - len(b_shape)
-    if diff == 0:
+
+    # Validate tensors have same rank and shape
+    if diff == 0 and all(a_shape[i] == b_shape[i] for i in range(len(a_shape))):
         return True
+
+    # Left-pad the shorter dimension with ones
     if diff > 0:
-        max = len(a_shape)
-        min = len(b_shape)
-        greater_tensor = a_shape
-        lesser_tensor = b_shape
-    elif diff < 0:
-        max = len(b_shape)
-        min = len(a_shape)
-        greater_tensor = b_shape
-        lesser_tensor = a_shape
-    j = min - 1
-    for i in range(max - 1, diff - 1, -1):
-        if not (
-            greater_tensor[i] != lesser_tensor[j]
-            and (greater_tensor[i] == 1 or lesser_tensor[i] == 1)
-        ):
+        b_shape = (1,) * abs(diff) + b_shape
+    else:
+        a_shape = (1,) * abs(diff) + a_shape
+
+    # Validate one of the following conditions for broadcastability per-dimension
+    # 1. Equal number of dimensions or 2. Dimension has shape 1
+    for i in range(len(a_shape)):
+        if not (a_shape[i] == b_shape[i] or a_shape[i] == 1 or b_shape[i] == 1):
             return False
     return True


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
@@ -1,8 +1,8 @@
 import torch
 import torch.nn as nn
+from harness import DispatchTestCase
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
-from harness import DispatchTestCase
 
 
 class TestWhereConverter(DispatchTestCase):
@@ -28,6 +28,20 @@ def forward(self, condition, x, y):
             expected_ops={torch.ops.aten.where.self},
         )
 
+    def test_0D_input(self):
+        class Where(nn.Module):
+            def forward(self, condition, x, y):
+                return torch.where(condition, x, y)
+
+        inputX = torch.randn((5, 6, 7, 1, 3))
+        inputOther = torch.tensor(8.0, dtype=torch.float)
+        condition = inputX < 0
+        self.run_test(
+            Where(),
+            (condition, inputX, inputOther),
+            expected_ops={torch.ops.aten.where.self},
+        )
+
 
 if __name__ == "__main__":
     run_tests()