chore: rebase and updates

Author: keehyuna

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -15,9 +15,11 @@
     ConverterRegistry,
     DynamoConverterImplSignature,
 )
-from torch_tensorrt.fx.converters.converter_utils import (
+from torch_tensorrt.fx.converters.converter_utils import (  # noqa: F401
     broadcast,
     get_axes_for_reduce_op,
+    prepend_ones,
+    set_layer_name,
 )
 from torch_tensorrt.fx.types import TRTDataType, TRTTensor
 

py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py

@@ -8,8 +8,9 @@
 from torch_tensorrt.dynamo.conversion.converter_utils import (
     broadcastable,
     get_trt_tensor,
+    prepend_ones,
+    set_layer_name,
 )
-from torch_tensorrt.fx.converters.converter_utils import prepend_ones, set_layer_name
 from torch_tensorrt.fx.types import TRTTensor
 
 

tests/py/dynamo/conversion/test_where_aten.py

@@ -77,59 +77,44 @@ def forward(self, condition):
             (condition,),
         )
 
-    # shape, min shape range, opt shape range, max shape range
+    # min/opt/max shape for condition/x/y input
     @parameterized.expand(
         [
             (
                 "3d_condition_3d_xshape_3d_yshape",
-                (-1, -1, -1),
                 (1, 1, 1),
                 (1, 2, 3),
                 (3, 3, 3),
-                (-1, -1, -1),
                 (1, 1, 1),
                 (1, 2, 3),
                 (3, 3, 3),
-                (-1, -1, -1),
                 (1, 1, 1),
                 (1, 2, 3),
                 (3, 3, 3),
             ),
             (
                 "1d_condition_3d_xshape_2d_yshape",
-                (-1),
                 (1,),
                 (2,),
                 (4,),
-                (-1, -1, -1),
                 (1, 1, 1),
                 (3, 2, 2),
                 (3, 2, 4),
-                (
-                    -1,
-                    -1,
-                ),
                 (1, 1),
                 (2, 2),
                 (2, 4),
             ),
             (
                 "2d_condition_3d_xshape_2d_yshape",
-                (-1, -1),
                 (4, 1),
                 (4, 2),
                 (5, 4),
-                (-1, -1, -1),
                 (1, 1, 1),
                 (3, 1, 2),
                 (3, 1, 4),
-                (
-                    -1,
-                    -1,
-                ),
-                (4, 1),
-                (4, 2),
-                (5, 4),
+                (1, 1),
+                (1, 2),
+                (1, 4),
             ),
         ]
     )
@@ -140,19 +125,22 @@ def forward(self, condition, x, y):
 
         input_specs = [
             Input(
-                shape=args[1],
+                min_shape=args[1],
+                opt_shape=args[2],
+                max_shape=args[3],
                 dtype=torch.bool,
-                shape_ranges=[(args[2], args[3], args[4])],
             ),
             Input(
-                shape=args[5],
+                min_shape=args[4],
+                opt_shape=args[5],
+                max_shape=args[6],
                 dtype=torch.float32,
-                shape_ranges=[(args[6], args[7], args[8])],
             ),
             Input(
-                shape=args[9],
+                min_shape=args[7],
+                opt_shape=args[8],
+                max_shape=args[9],
                 dtype=torch.float32,
-                shape_ranges=[(args[10], args[11], args[12])],
             ),
         ]
         self.run_test_with_dynamic_shape(Where(), input_specs)