Fixing the acc tests, logical_and operator and the leaky_relu test

Author: apbose

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -1262,12 +1262,7 @@ def acc_ops_logical_not(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_val = kwargs["input"]
-    operation_type = trt.UnaryOperation.NOT
-    # cast to bool type
-    if input_val.dtype in (trt.float32, trt.float16, trt.int32):
-        input_val = type_cast(network, target, f"{name}_input", input_val, trt.bool)
-    return add_unary_layer(network, input_val, operation_type, target, name)
+    return add_logical_not(network, target, kwargs, name)
 
 
 @tensorrt_converter(acc_ops.logical_and, no_implicit_batch_dim=True)
@@ -2335,7 +2330,7 @@ def acc_ops_getitem(
     input_val = kwargs["input"]
     slices = kwargs["idx"]
     if not isinstance(input_val, TRTTensor):
-        return getitem(input_val, slices)  # type: ignore[arg-type]
+        return operator.getitem(input_val, slices)  # type: ignore[arg-type]
 
     if not isinstance(slices, tuple) and not isinstance(slices, list):
         slices = (slices,)
@@ -2803,7 +2798,7 @@ def acc_ops_hardtanh(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return add_hardtanh(network, target, kwargs, name)
+    return add_hard_tanh(network, target, kwargs, name)
 
 
 @tensorrt_converter(acc_ops.interpolate)

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -489,9 +489,7 @@ def aten_ops_leaky_relu(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    kwargs_new = {
-        "input": args[0],
-    }
+    kwargs_new = {"input": args[0], "negative_slope": args[1]}
     return add_leaky_relu(network, target, kwargs_new, name)
 
 

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -328,63 +328,6 @@ def broadcast(
     return a, b
 
 
-def get_shape_with_dynamic_shape(
-    network: TRTNetwork,
-    shape: Union[list, tuple, torch.Tensor],
-    input_val: TRTTensor,
-    target: Target,
-    name: str,
-) -> TRTTensor:
-    """
-    Prepare the real output tensor shape for dynamic shape mode tensor input.
-    How this functions works:
-    Assuming the input_val has actual shape [2048, 256, 512], expected reduce operation
-    output shape is [-1, 128, 256], this function should return [2048, 128, 256] as the actual
-    reduce operation output shape. Steps of calculations are:
-        1. get the actual tensor shape of input_val via add_shape layer;
-        2. create a all 0 tensor [0, 0, 0];
-        3. run elementwise comparision the [0, 0, 0] and [-1, 128, 256] tensor, get a condition tensor [True, False, False];
-        4. use the condition tensor [True, False, False] to do selection between [2048, 256, 512] and [-1, 128, 256], replace
-           all -1 dynamic shape dimensions with actual batch_size value;
-        5. output shape with actual batch_size as [2048, 128, 256]
-
-    Args:
-        network (TRTNetwork): TensorRT network object.
-        shape: calculated shape of the expected output tensor
-        input_val (TRTTensor): A TensorRT ITensor.
-        target (Target): Target of fx node.
-        name (str): The name we want to assign to the created TensorRT layer.
-    Returns:
-        TensorRT ITensors that represents the actual shape of the input_val
-    """
-    # Ger real shape info for input_val
-    input_shape = network.add_shape(input_val).get_output(0)
-
-    scale_layer = network.add_constant(
-        input_shape.shape, np.ascontiguousarray(shape, dtype=np.int32)
-    )
-    set_layer_name(scale_layer, target, f"{name}_scale")
-    scale_res = scale_layer.get_output(0)
-
-    length = input_shape.shape[0]
-    zero_layer = network.add_constant(
-        input_shape.shape, to_numpy(torch.zeros((length), dtype=torch.int32))
-    )
-    set_layer_name(zero_layer, target, f"{name}_zeros")
-
-    condition_val = operator.add_binary_elementwise_layer(
-        network,
-        scale_res,
-        zero_layer.get_output(0),
-        trt.ElementWiseOperation.LESS,
-        target,
-        f"{name}_shape",
-    )
-    select_layer = network.add_select(condition_val, input_shape, scale_res)
-    set_layer_name(select_layer, target, f"{name}_select")
-    return select_layer.get_output(0)
-
-
 def squeeze_left(const: torch.Tensor):
     """
     Squeeze the size-1 dimensions on the left side of the shape tuple.
@@ -529,22 +472,6 @@ def dtype_uniform(
     return input, other
 
 
-def type_cast(
-    network: TRTNetwork,
-    target: Target,
-    name: str,
-    input: TRTTensor,
-    cast_type: TRTDataType,
-):
-    """
-    This function helps to cast the input type to cast_type
-    """
-    layer_i = network.add_identity(input)
-    layer_i.set_output_type(0, cast_type)
-    set_layer_name(layer_i, target, f"{name}_dtype_change")
-    return layer_i.get_output(0)
-
-
 def to_numpy(tensor: Optional[torch.Tensor]) -> Optional[np.ndarray]:
     """
     Convert a PyTorch Tensor to a Numpy Array. If the tensor is

py/torch_tensorrt/fx/converters/operator.py

@@ -21,7 +21,6 @@
 from .converter_utils import get_positive_dim
 from .converter_utils import prepend_ones
 from .converter_utils import has_dynamic_shape
-from .converter_utils import get_shape_with_dynamic_shape
 from .converter_utils import to_numpy
 
 from ..types import (
@@ -289,6 +288,79 @@ def trt_dtype_to_torch_dtype(trt_dtype):
     return table[trt_dtype]
 
 
+def get_shape_with_dynamic_shape(
+    network: TRTNetwork,
+    shape: Union[list, tuple, torch.Tensor],
+    input_val: TRTTensor,
+    target: Target,
+    name: str,
+) -> TRTTensor:
+    """
+    Prepare the real output tensor shape for dynamic shape mode tensor input.
+    How this functions works:
+    Assuming the input_val has actual shape [2048, 256, 512], expected reduce operation
+    output shape is [-1, 128, 256], this function should return [2048, 128, 256] as the actual
+    reduce operation output shape. Steps of calculations are:
+        1. get the actual tensor shape of input_val via add_shape layer;
+        2. create a all 0 tensor [0, 0, 0];
+        3. run elementwise comparision the [0, 0, 0] and [-1, 128, 256] tensor, get a condition tensor [True, False, False];
+        4. use the condition tensor [True, False, False] to do selection between [2048, 256, 512] and [-1, 128, 256], replace
+           all -1 dynamic shape dimensions with actual batch_size value;
+        5. output shape with actual batch_size as [2048, 128, 256]
+
+    Args:
+        network (TRTNetwork): TensorRT network object.
+        shape: calculated shape of the expected output tensor
+        input_val (TRTTensor): A TensorRT ITensor.
+        target (Target): Target of fx node.
+        name (str): The name we want to assign to the created TensorRT layer.
+    Returns:
+        TensorRT ITensors that represents the actual shape of the input_val
+    """
+    # Ger real shape info for input_val
+    input_shape = network.add_shape(input_val).get_output(0)
+
+    scale_layer = network.add_constant(
+        input_shape.shape, np.ascontiguousarray(shape, dtype=np.int32)
+    )
+    set_layer_name(scale_layer, target, f"{name}_scale")
+    scale_res = scale_layer.get_output(0)
+
+    length = input_shape.shape[0]
+    zero_layer = network.add_constant(
+        input_shape.shape, to_numpy(torch.zeros((length), dtype=torch.int32))
+    )
+    set_layer_name(zero_layer, target, f"{name}_zeros")
+
+    condition_val = add_binary_elementwise_layer(
+        network,
+        scale_res,
+        zero_layer.get_output(0),
+        trt.ElementWiseOperation.LESS,
+        target,
+        f"{name}_shape",
+    )
+    select_layer = network.add_select(condition_val, input_shape, scale_res)
+    set_layer_name(select_layer, target, f"{name}_select")
+    return select_layer.get_output(0)
+
+
+def type_cast(
+    network: TRTNetwork,
+    target: Target,
+    name: str,
+    input: TRTTensor,
+    cast_type: TRTDataType,
+):
+    """
+    This function helps to cast the input type to cast_type
+    """
+    layer_i = network.add_identity(input)
+    layer_i.set_output_type(0, cast_type)
+    set_layer_name(layer_i, target, f"{name}_dtype_change")
+    return layer_i.get_output(0)
+
+
 def add_tile(network, target, kwargs, name):
     input_t = kwargs["input"]
     input_val = get_trt_tensor(network, input_t, f"{name}_input")
@@ -822,25 +894,54 @@ def add_minimum(network, target, kwargs, name):
     )
 
 
+def add_logical_not(network, target, kwargs, name):
+    input_val = kwargs["input"]
+    operation_type = trt.UnaryOperation.NOT
+    # cast to bool type
+    if input_val.dtype in (trt.float32, trt.float16, trt.int32):
+        input_val = type_cast(network, target, f"{name}_input", input_val, trt.bool)
+    return add_unary_layer(network, input_val, operation_type, target, name)
+
+
 def add_logical_and(network, target, kwargs, name):
     if network.has_implicit_batch_dimension:
         raise RuntimeError(
-            "The `ne` function should be called with explicit batch dimension."
+            "The `logical_and` function should be called with explicit batch dimension."
         )
 
     input_t = kwargs["input"]
     other_t = kwargs["other"]
+    # we only support both inputs are bool type
+    if target == acc_ops.bitwise_and:
+
+        def check_is_bool(input_t):
+            if isinstance(input_t, TRTTensor):
+                assert (
+                    input_t.dtype == trt.bool
+                ), "We currently do not support input is non-bool"
+            elif isinstance(input_t, torch.Tensor):
+                assert (
+                    input_t.dtype == torch.bool
+                ), "We currently do not support input is non-bool"
+            else:
+                assert isinstance(
+                    input_t.bool
+                ), "We currently do not support input is non-bool"
+
+        check_is_bool(input_t)
+        check_is_bool(other_t)
 
     input_t = get_trt_tensor(network, input_t, f"{name}_input_t")
     other_t = get_trt_tensor(network, other_t, f"{name}_other_t")
 
-    input_t, other_t = dtype_uniform(network, target, name, input_t, other_t)
-    eq_t = add_binary_elementwise_layer(
-        network, input_t, other_t, trt.ElementWiseOperation.EQUAL, target, name
+    if input_t.dtype != trt.bool:
+        input_t = type_cast(network, target, f"{name}_input", input_t, trt.bool)
+    if other_t.dtype != trt.bool:
+        other_t = type_cast(network, target, f"{name}_other", other_t, trt.bool)
+    return add_binary_elementwise_layer(
+        network, input_t, other_t, trt.ElementWiseOperation.AND, target, name
     )
 
-    return add_unary_layer(network, eq_t, trt.UnaryOperation.NOT, target, name)
-
 
 def add_ne(network, target, kwargs, name):
     if network.has_implicit_batch_dimension:

py/torch_tensorrt/fx/test/converters/aten_op/test_leaky_relu_aten.py

@@ -8,7 +8,7 @@ class TestLeakyReLUConverter(DispatchTestCase):
     def test_leaky_relu(self):
         class TestModule(nn.Module):
             def forward(self, x):
-                return nn.functional.leaky_relu(x)
+                return nn.functional.leaky_relu(x, negative_slope=0.05)
 
         inputs = [torch.randn(1, 10)]
         self.run_test(
@@ -18,7 +18,7 @@ def forward(self, x):
     def test_leaky_relu_with_dynamic_shape(self):
         class TestModule(nn.Module):
             def forward(self, x):
-                return nn.functional.leaky_relu(x)
+                return nn.functional.leaky_relu(x, negative_slope=0.05)
 
         input_specs = [
             InputTensorSpec(
@@ -34,7 +34,7 @@ def forward(self, x):
     def test_leaky_relu_with_dynamic_shape_four_dimensions(self):
         class TestModule(nn.Module):
             def forward(self, x):
-                return nn.functional.leaky_relu(x)
+                return nn.functional.leaky_relu(x, negative_slope=0.05)
 
         input_specs = [
             InputTensorSpec(