chore: move threshold values to variables and rebase with main

Author: peri044

py/torch_tensorrt/dynamo/utils.py

@@ -22,6 +22,8 @@
 
 COSINE_THRESHOLD = 0.99
 DYNAMIC_DIM = -1
+RTOL = 5e-3
+ATOL = 5e-3
 
 
 class Frameworks(Enum):
@@ -412,6 +414,8 @@ def check_module_output(
 def check_output_equal(
     output1: Any,
     output2: Any,
+    rtol: float = RTOL,
+    atol: float = ATOL,
 ) -> bool:
 
     if type(output1) != type(output2):
@@ -423,7 +427,7 @@ def check_output_equal(
     if isinstance(output1, torch.Tensor):
         if output1.shape != output2.shape:
             return False
-        return torch.allclose(output1, output2, 5e-3, 5e-3)  # type: ignore
+        return torch.allclose(output1, output2, rtol, atol)  # type: ignore
 
     elif isinstance(output1, (tuple, list)):
         if len(output1) != len(output2):

tests/py/dynamo/conversion/harness.py

@@ -23,7 +23,7 @@
     pre_export_lowering,
 )
 from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule
-from torch_tensorrt.dynamo.utils import get_torch_inputs
+from torch_tensorrt.dynamo.utils import ATOL, RTOL, get_torch_inputs
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -60,8 +60,8 @@ def run_test(
         mod,
         inputs,
         interpreter,
-        rtol,
-        atol,
+        rtol=RTOL,
+        atol=ATOL,
         check_dtype=True,
         pyt_inputs=None,
         rt_cls=PythonTorchTensorRTModule,
@@ -254,8 +254,8 @@ def run_test(
         self,
         mod,
         inputs,
-        rtol=5e-3,
-        atol=5e-3,
+        rtol=RTOL,
+        atol=ATOL,
         precision=dtype.f32,
         check_dtype=True,
         use_dynamo_tracer=False,
@@ -374,8 +374,8 @@ def run_test_with_dynamic_shape(
         self,
         mod,
         input_specs,
-        rtol=5e-3,
-        atol=5e-3,
+        rtol=RTOL,
+        atol=ATOL,
         output_dtypes=None,
         use_dynamo_tracer=False,
         enable_passes=False,

tests/py/dynamo/conversion/test_bitwise_and_aten.py

@@ -5,6 +5,7 @@
 from torch.export import Dim
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
+from torch_tensorrt.dynamo.utils import ATOL, RTOL
 
 from .harness import DispatchTestCase
 
@@ -152,8 +153,8 @@ def forward(self, lhs_val, rhs_val):
                 torch.testing.assert_close(
                     out,
                     ref,
-                    rtol=5e-3,
-                    atol=5e-3,
+                    rtol=RTOL,
+                    atol=ATOL,
                     equal_nan=True,
                     check_dtype=True,
                 )

tests/py/dynamo/conversion/test_embedding_bag_aten.py

@@ -3,6 +3,7 @@
 from parameterized import param, parameterized
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
+from torch_tensorrt.dynamo.utils import ATOL, RTOL
 
 from .harness import DispatchTestCase
 
@@ -501,8 +502,8 @@ def forward(self, weights, indices, offsets, per_sample_weights=None):
                 torch.testing.assert_close(
                     out,
                     ref,
-                    rtol=5e-3,
-                    atol=5e-3,
+                    rtol=RTOL,
+                    atol=ATOL,
                     equal_nan=True,
                     check_dtype=True,
                 )

tests/py/dynamo/conversion/test_index_select_aten.py

@@ -4,6 +4,7 @@
 from parameterized import param, parameterized
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
+from torch_tensorrt.dynamo.utils import ATOL, RTOL
 
 from .harness import DispatchTestCase
 
@@ -122,8 +123,8 @@ def forward(self, source_tensor, indice_tensor):
                 torch.testing.assert_close(
                     out,
                     ref,
-                    rtol=5e-3,
-                    atol=5e-3,
+                    rtol=RTOL,
+                    atol=ATOL,
                     equal_nan=True,
                     check_dtype=True,
                 )