update test case

Author: bowang007

py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin.py

@@ -130,23 +130,25 @@ def _generic_plugin_desc(*args: Any, **kwargs: Any) -> Tuple[trtp.TensorDesc]:
             output = torch_op(*fake_args, **kwargs)
 
         # We assume that number of dimensions are the same in torch op
-        # shape_calc_fns = [None] * args[0].ndim
         shape_calc_fns = [None] * output.ndim
-        
+
         for i in range(output.ndim):
-            input_node_expr = [syms_arg[j].node.expr for syms_arg in syms_args for j in range(len(syms_arg))]
-            shape_calc_fns[i] = lambdify(tuple(input_node_expr), output.shape[i].node.expr, "math")
-        
-        # for i in range(args[0].ndim):
-        #     input_node_expr = [syms_arg[i].node.expr for syms_arg in syms_args]
-        #     shape_calc_fns[i] = lambdify(
-        #         tuple(input_node_expr), output.shape[i].node.expr, "math"
-        #     )
+            input_node_expr = [
+                syms_arg[j].node.expr
+                for syms_arg in syms_args
+                for j in range(len(syms_arg))
+            ]
+            shape_calc_fns[i] = lambdify(
+                tuple(input_node_expr), output.shape[i].node.expr, "math"
+            )
 
         out_desc = tensor_args[0].like()
         for i in range(out_desc.ndim):
-            input_shape_expr = [arg.shape_expr[j] for arg in tensor_args for j in range(len(arg.shape_expr))]
-            # input_shape_expr = [tensor_arg.shape_expr[i] for tensor_arg in tensor_args]
+            input_shape_expr = [
+                arg.shape_expr[j]
+                for arg in tensor_args
+                for j in range(len(arg.shape_expr))
+            ]
             if output.shape[i].node.expr is None:
                 raise ValueError(f"output.shape[{i}].node.expr cannot be None")
             out_desc.shape_expr[i] = shape_calc_fns[i](*input_shape_expr)  # type: ignore[misc]

tests/py/dynamo/automatic_plugin/test_automatic_plugin.py

@@ -81,12 +81,3 @@ def forward(self, lhs, rhs):
 
 if __name__ == "__main__":
     run_tests()
-
-# Example Usage
-# A = torch.full((64, 64), 2, device="cuda", dtype=torch.float)
-# B = torch.full((64, 64), 3, device="cuda", dtype=torch.float)
-
-# C, D = torch.ops.torchtrt_ex.elementwise_add_mul.default(A, B)
-
-# print("C (Addition):", C)
-# print("D (Multiplication):", D)

tests/py/dynamo/automatic_plugin/test_flashinfer_rmsnorm.py

@@ -1,49 +1,50 @@
 import flashinfer
-
 import torch
 import torch.nn as nn
 import torch_tensorrt
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt._enums import dtype
 
 from ..conversion.harness import DispatchTestCase
-import flashinfer
 
 
-@torch.library.custom_op("torchtrt_ex::flashinfer_rmsnorm", mutates_args=())  # type: ignore[misc]
+@torch.library.custom_op("flashinfer::rmsnorm", mutates_args=())  # type: ignore[misc]
 def flashinfer_rmsnorm(
     input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
 ) -> torch.Tensor:
     return flashinfer.norm.rmsnorm(input, weight)
 
 
-@torch.library.register_fake("torchtrt_ex::flashinfer_rmsnorm")
+@torch.library.register_fake("flashinfer::rmsnorm")
 def _(input: torch.Tensor, weight: torch.Tensor, b: float = 1e-6) -> torch.Tensor:
     return input
 
 
-
 torch_tensorrt.dynamo.conversion.plugins.custom_op(
-    "torchtrt_ex::flashinfer_rmsnorm", supports_dynamic_shapes=True
+    "flashinfer::rmsnorm", supports_dynamic_shapes=True
 )
 
 
 class TestAutomaticPlugin(DispatchTestCase):
     @parameterized.expand(
         [
-            ((64, 64), (64, ), torch.float16),
-            ((256, 256), (256, ), torch.float16),
+            ((64, 64), (64,), torch.float16),
+            ((256, 256), (256,), torch.float16),
         ]
     )
-    def test_rmsnorm_float(self, input_shape, weight_shape, dtype):
+    def test_rmsnorm_float(self, input_shape, weight_shape, data_type):
         class rmsnorm(nn.Module):
             def forward(self, input, weight):
-                return torch.ops.torchtrt_ex.flashinfer_rmsnorm.default(input, weight)
+                return torch.ops.flashinfer.rmsnorm.default(input, weight)
 
-        inputs = [torch.randn(input_shape, device="cuda", dtype=dtype),  torch.randn(weight_shape, device="cuda", dtype=dtype)]
+        inputs = [
+            torch.randn(input_shape, device="cuda", dtype=data_type),
+            torch.randn(weight_shape, device="cuda", dtype=data_type),
+        ]
 
-        self.run_test(rmsnorm(), inputs)
+        self.run_test(rmsnorm(), inputs, precision=dtype.f16)
 
 
 if __name__ == "__main__":
-    run_tests()
+    run_tests()