Add dynamic shape suport for acc_ops.max_pool1d (#83)

Summary:
Pull Request resolved: https://github.com/pytorch/fx2trt/pull/83

Initial draft for adding dynamic shape support for acc_ops.max_pool1d.

Reviewed By: frank-wei, wushirong

Differential Revision: D36617147

fbshipit-source-id: 162cf66e0c8539f74a6c009e8eee624038b50d13

Author: ShreyanshPrajapati

fx/converters/acc_ops_converters.py

@@ -1849,21 +1849,21 @@ def acc_ops_embedding(
     return gather_layer.get_output(0)
 
 
-@tensorrt_converter(acc_ops.max_pool1d, no_explicit_batch_dim=True)
+@tensorrt_converter(acc_ops.max_pool1d)
 def acc_ops_max_pool1d(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    if not network.has_implicit_batch_dimension:
+    input_trt = kwargs["input"]
+    if not isinstance(input_trt, TRTTensor):
         raise RuntimeError(
-            "Current implementation does not support dynamic shape. Make sure that the network has an explicit batch dimension!"
+            f"Max_pool1d received input {input_trt} that is not part "
+            "of the TensorRT region!"
         )
 
-    input_trt = kwargs["input"]
-
     # adds unsqueeze layer -> max pool 2d -> squeeze layer to emulate max pool 1d.
     unsqueeze_layer = network.add_shuffle(input=input_trt)
     unsqueeze_layer.reshape_dims = tuple([*input_trt.shape, 1])

test/converters/acc_op/test_maxpool.py

@@ -40,7 +40,32 @@ def forward(self, x):
             TestModule(),
             inputs,
             expected_ops={acc_ops.max_pool1d},
-            test_explicit_batch_dim=False,
+        )
+
+    def test_max_pool1d_with_dynamic_shape(
+        self,
+    ):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.max_pool = torch.nn.MaxPool1d(1)
+
+            def forward(self, x):
+                return self.max_pool(x)
+
+        # shape is not set to (-1, -1, -1) as reshape dimension with
+        # more than one -1 wildcard is not allowed while adding unsqueeze layer
+        input_specs = [
+            InputTensorSpec(
+                shape=(1, 1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 1, 4), (1, 1, 4))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(),
+            input_specs,
+            expected_ops={acc_ops.max_pool1d},
         )
 
     @parameterized.expand(