Add dynamic shape support for leaky_relu/elu/hard_sigmoid/softplus (#2927)

Author: lanluo-nvidia

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -459,7 +459,9 @@ def aten_ops_tanh(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.leaky_relu.default)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.leaky_relu.default, supports_dynamic_shapes=True
+)
 def aten_ops_leaky_relu(
     ctx: ConversionContext,
     target: Target,
@@ -477,7 +479,7 @@ def aten_ops_leaky_relu(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.elu.default)
+@dynamo_tensorrt_converter(torch.ops.aten.elu.default, supports_dynamic_shapes=True)
 def aten_ops_elu(
     ctx: ConversionContext,
     target: Target,
@@ -496,7 +498,9 @@ def aten_ops_elu(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.softplus.default)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.softplus.default, supports_dynamic_shapes=True
+)
 def aten_ops_softplus(
     ctx: ConversionContext,
     target: Target,
@@ -514,7 +518,9 @@ def aten_ops_softplus(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.hardsigmoid.default)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.hardsigmoid.default, supports_dynamic_shapes=True
+)
 def aten_ops_hard_sigmoid(
     ctx: ConversionContext,
     target: Target,

tests/py/dynamo/conversion/test_elu_aten.py

@@ -22,9 +22,10 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, -1, -1),
                 dtype=torch.float32,
-                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+                min_shape=(1, 1, 1),
+                opt_shape=(1, 2, 3),
+                max_shape=(3, 3, 3),
             ),
         ]
         self.run_test_with_dynamic_shape(TestModule(), input_specs)
@@ -36,9 +37,10 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, -1, -1, -1),
                 dtype=torch.float32,
-                shape_ranges=[((1, 1, 1, 5), (1, 2, 3, 5), (3, 3, 3, 5))],
+                min_shape=(1, 1, 1, 5),
+                opt_shape=(1, 2, 3, 5),
+                max_shape=(3, 3, 3, 5),
             ),
         ]
 

tests/py/dynamo/conversion/test_hard_sigmoid_aten.py

@@ -22,9 +22,10 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, -1, -1),
                 dtype=torch.float32,
-                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+                min_shape=(1, 1, 1),
+                opt_shape=(1, 2, 3),
+                max_shape=(3, 3, 3),
             ),
         ]
         self.run_test_with_dynamic_shape(TestModule(), input_specs)
@@ -36,9 +37,10 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, -1, -1, -1),
                 dtype=torch.float32,
-                shape_ranges=[((1, 1, 1, 5), (1, 2, 3, 5), (3, 3, 3, 5))],
+                min_shape=(1, 1, 1, 5),
+                opt_shape=(1, 2, 3, 5),
+                max_shape=(3, 3, 3, 5),
             ),
         ]
 

tests/py/dynamo/conversion/test_leaky_relu_aten.py

@@ -22,9 +22,10 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, -1, -1),
                 dtype=torch.float32,
-                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+                min_shape=(1, 1, 1),
+                opt_shape=(1, 2, 3),
+                max_shape=(3, 3, 3),
             ),
         ]
         self.run_test_with_dynamic_shape(TestModule(), input_specs)
@@ -36,9 +37,10 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, -1, -1, -1),
                 dtype=torch.float32,
-                shape_ranges=[((1, 1, 1, 5), (1, 2, 3, 5), (3, 3, 3, 5))],
+                min_shape=(1, 1, 1, 5),
+                opt_shape=(1, 2, 3, 5),
+                max_shape=(3, 3, 3, 5),
             ),
         ]
 

tests/py/dynamo/conversion/test_softplus_aten.py

@@ -22,9 +22,10 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, -1, -1),
                 dtype=torch.float32,
-                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+                min_shape=(1, 1, 1),
+                opt_shape=(1, 2, 3),
+                max_shape=(3, 3, 3),
             ),
         ]
         self.run_test_with_dynamic_shape(TestModule(), input_specs)
@@ -36,9 +37,10 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, -1, -1, -1),
                 dtype=torch.float32,
-                shape_ranges=[((1, 1, 1, 5), (1, 2, 3, 5), (3, 3, 3, 5))],
+                min_shape=(1, 1, 1, 5),
+                opt_shape=(1, 2, 3, 5),
+                max_shape=(3, 3, 3, 5),
             ),
         ]