Converter reorg and rsub

Rsub error fixes and linting error fixed

Rsub test case to include different inputs

Author: apbose

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -24,6 +24,7 @@
 from torch_tensorrt.fx.converters.impl.elementwise import trunc_div
 from torch_tensorrt.fx.converters.impl.elementwise import rsqrt
 from torch_tensorrt.fx.converters.impl.elementwise import fmod
+from torch_tensorrt.fx.converters.impl.elementwise import rsub
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -452,6 +453,20 @@ def aten_ops_reshape(
     return layer.get_output(0)
 
 
+@tensorrt_converter(torch.ops.aten.rsub.Tensor)
+def aten_ops_rsub(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    alpha = None
+    if "alpha" in kwargs:
+        alpha = kwargs["alpha"]
+    return rsub(network, target, SourceIR.ATEN, name, args[0], args[1], alpha)
+
+
 @tensorrt_converter(torch.ops.aten.tanh.default)
 def aten_ops_tanh(
     network: TRTNetwork,
@@ -460,7 +475,6 @@ def aten_ops_tanh(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-
     return activation.tanh(
         network,
         target,

py/torch_tensorrt/fx/converters/impl/elementwise/ops.py

@@ -141,6 +141,44 @@ def rsqrt(
     return output
 
 
+def rsub(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    other: TRTTensor,
+) -> TRTTensor:
+    # NOTE: TRT doesnt currently implement fmod so we need multiple operations to perform it
+    trunc_div_value = trunc_div(
+        network,
+        target,
+        source_ir,
+        name + "_trunc_div",
+        input,
+        other,
+    )
+    prod_value = convert_binary_elementwise(
+        network,
+        target,
+        source_ir,
+        name + "_prod",
+        trt.ElementWiseOperation.PROD,
+        trunc_div_value,
+        other,
+    )
+    sub_value = convert_binary_elementwise(
+        network,
+        target,
+        SourceIR.ACC,
+        name + "_sub",
+        trt.ElementWiseOperation.SUB,
+        input,
+        prod_value,
+    )
+    return sub_value
+
+
 def fmod(
     network: TRTNetwork,
     target: Target,

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

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests, TestCase
+from torch_tensorrt.dynamo import compile
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestRSubConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("2d_dim_alpha", (2, 1), 2),
+            ("3d_dim_alpha", (2, 1, 2), 2),
+        ]
+    )
+    def test_rsub_same(self, _, x, alpha):
+        class rsub(nn.Module):
+            def forward(self, input):
+                return torch.rsub(input, input, alpha=alpha)
+
+        inputs = [torch.randn(x)]
+        self.run_test(
+            rsub(),
+            inputs,
+            expected_ops={torch.ops.aten.rsub.Tensor},
+        )
+
+    # @parameterized.expand(
+    #     [
+    #         ("2d_dim_alpha", (2, 1), 2),
+    #         ("3d_dim_alpha", (2, 1, 2), 2),
+    #     ]
+    # )
+    # def test_rsub_diff(self, _, x, alpha):
+    #     class rsub(nn.Module):
+    #         def forward(self, inputOne, inputTwo):
+    #             return torch.rsub(inputOne, inputTwo, alpha=alpha)
+
+    #     inputOne = [torch.randn(x)]
+    #     inputTwo = [torch.randn(x)]
+    #     inputs = (inputOne, inputTwo)
+    #     self.run_test(
+    #         rsub(),
+    #         inputs,
+    #         expected_ops={torch.ops.aten.rsub.Tensor},
+    #     )
+
+class TestRSubDiff(TestCase):
+    def test_rsub_diff(self):
+        class rsub_diff(nn.Module):
+            def forward(self, inputOne, inputTwo):
+                return torch.rsub(inputOne, inputTwo, alpha=2)
+        inputOne = torch.randn(2,1).cuda()
+        inputTwo = torch.randn(2,1).cuda()
+        alpha = 2
+        inputs = [inputOne, inputTwo]
+        fx_graph = torch.fx.symbolic_trace(rsub_diff())
+        torch._dynamo.reset()
+        optimized_model = compile(
+            fx_graph, inputs, min_block_size=1, pass_through_build_failures=True
+        )
+        optimized_model_results = optimized_model(*inputs).detach().cpu()
+        torch_model_results = fx_graph(*inputs).detach().cpu()
+        max_diff = float(
+            torch.max(torch.abs(optimized_model_results - torch_model_results))
+        )
+        self.assertAlmostEqual(
+            max_diff,
+            0,
+            5,
+            f"Reciprocal TRT outputs don't match with the original model.",
+        )
+
+if __name__ == "__main__":
+    run_tests()
+    #Test two
+    
+    
+    
+