No actual change, just remove variable contain Tensors from global scope (pytorch#143225)

Pull Request resolved: pytorch#143225
Approved by: https://github.com/ezyang

Author: albanD

test/jit/test_complexity.py

@@ -20,8 +20,8 @@
 from torch.testing._internal.jit_metaprogramming_utils import (
     get_all_nn_module_tests,
     get_nn_functional_compiled_fn_and_inputs,
+    get_nn_functional_tests,
     get_nn_mod_test_name,
-    nn_functional_tests,
     try_get_nn_module_compiled_mod_and_inputs,
 )
 from torch.testing._internal.jit_utils import enable_profiling_mode, JitTestCase
@@ -70,7 +70,7 @@ def tearDown(self):
     def test_generated_functional_tests(self):
         with enable_profiling_mode():
             stats = [("Name", "Ifs/Loops", "non-tensor ops")]
-            for test in nn_functional_tests:
+            for test in get_nn_functional_tests():
                 test_name = test[0]
 
                 fn, inputs = get_nn_functional_compiled_fn_and_inputs(*test)

test/test_cpp_api_parity.py

@@ -42,7 +42,7 @@ class TestCppApiParity(common.TestCase):
         (sample_module.module_tests, common_nn.NewModuleTest),
         (sample_functional.functional_tests, common_nn.NewModuleTest),
         (common_nn.module_tests, common_nn.NewModuleTest),
-        (common_nn.new_module_tests, common_nn.NewModuleTest),
+        (common_nn.get_new_module_tests(), common_nn.NewModuleTest),
         (common_nn.criterion_tests, common_nn.CriterionTest),
     ]:
         for test_params_dict in test_params_dicts:

test/test_expanded_weights.py

@@ -25,7 +25,11 @@
 )
 from torch.testing._internal.common_methods_invocations import op_db, SampleInput
 from torch.testing._internal.common_modules import module_db, modules
-from torch.testing._internal.common_nn import module_tests, new_module_tests, TestBase
+from torch.testing._internal.common_nn import (
+    get_new_module_tests,
+    module_tests,
+    TestBase,
+)
 from torch.testing._internal.common_utils import (
     freeze_rng_state,
     make_tensor,
@@ -1011,7 +1015,7 @@ def filter_supported_tests(t):
 # TODO: Once all of these use ModuleInfo, replace with ModuleInfo tests
 # These currently use the legacy nn tests
 supported_tests = [
-    t for t in module_tests + new_module_tests if filter_supported_tests(t)
+    t for t in module_tests + get_new_module_tests() if filter_supported_tests(t)
 ]
 for test_param in supported_tests:
     if "constructor" not in test_param:

test/test_fx_experimental.py

@@ -50,7 +50,7 @@
     ops,
 )
 from torch.testing._internal.common_methods_invocations import op_db
-from torch.testing._internal.common_nn import module_tests, new_module_tests
+from torch.testing._internal.common_nn import module_tests, get_new_module_tests
 from torch.testing._internal.common_utils import TEST_Z3, run_tests, TestCase
 from torch.testing._internal.jit_utils import JitTestCase
 import torch.utils._pytree as pytree
@@ -1006,7 +1006,7 @@ def test_normalize_modules_exhaustive(self):
         Exhaustively test `Node.normalized_arguments` on all standard
         torch.nn Module classes
         """
-        for test_params in module_tests + new_module_tests:
+        for test_params in module_tests + get_new_module_tests():
             if "constructor" not in test_params:
                 constructor = getattr(torch.nn, test_params["module_name"])
             else:

test/test_jit.py

@@ -107,10 +107,10 @@
 from torch.testing._internal.jit_metaprogramming_utils import (
     get_script_args,
     create_input, unpack_variables,
-    additional_module_tests, EXCLUDE_SCRIPT_MODULES,
+    get_all_nn_module_tests, EXCLUDE_SCRIPT_MODULES,
     get_nn_module_name_from_kwargs, get_nn_mod_test_name, script_method_template)
 
-from torch.testing._internal.common_nn import module_tests, new_module_tests, criterion_tests
+from torch.testing._internal.common_nn import criterion_tests
 
 # For testing truediv in python 2
 from torch.testing._internal.test_module.future_div import div_int_future, div_float_future
@@ -16247,7 +16247,7 @@ def test_version(self):
         # issue gh-32561
         self.assertTrue(torch.__version__.startswith(torch.onnx.producer_version))
 
-for test in module_tests + new_module_tests + additional_module_tests:
+for test in get_all_nn_module_tests():
     add_nn_module_test(**test)
 
 for test in criterion_tests:

test/test_nn.py

@@ -38,7 +38,7 @@
 from torch.testing._internal.common_cuda import TEST_CUDA, TEST_MULTIGPU, TEST_CUDNN, PLATFORM_SUPPORTS_FLASH_ATTENTION
 from torch.testing._internal.common_nn import NNTestCase, NewModuleTest, CriterionTest, \
     module_tests, criterion_tests, loss_reference_fns, _create_basic_net, \
-    ctcloss_reference, new_module_tests, single_batch_reference_fn, _test_bfloat16_ops, _test_module_empty_input
+    ctcloss_reference, get_new_module_tests, single_batch_reference_fn, _test_bfloat16_ops, _test_module_empty_input
 from torch.testing._internal.common_device_type import dtypesIfMPS, instantiate_device_type_tests, dtypes, \
     dtypesIfCUDA, precisionOverride, skipCUDAIfCudnnVersionLessThan, onlyCUDA, onlyCPU, \
     skipCUDAIfRocm, skipCUDAIf, skipCUDAIfNotRocm, \
@@ -7332,7 +7332,7 @@ def with_tf32_on(self, test=test, kwargs=kwargs):
         else:
             add(cuda_test_name, with_tf32_off)
 
-for test_params in module_tests + new_module_tests:
+for test_params in module_tests + get_new_module_tests():
     # TODO: CUDA is not implemented yet
     if 'constructor' not in test_params:
         name = test_params.pop('module_name')

torch/ao/quantization/pt2e/qat_utils.py

@@ -19,8 +19,6 @@
 from torch.fx.subgraph_rewriter import replace_pattern_with_filters, ReplacedPatterns
 
 from .utils import (
-    _conv1d_bn_example_inputs,
-    _conv2d_bn_example_inputs,
     _get_aten_graph_module_for_pattern,
     _is_bn_node,
     _is_conv_or_conv_transpose_node,
@@ -35,27 +33,6 @@
 __all__ = []  # type: ignore[var-annotated]
 
 
-# Example inputs for quantized and folded conv-bn1d patterns used in convert
-_quantized_conv1d_bn_example_inputs = (
-    torch.randn(1, 1, 3),  # x
-    torch.randn(1, 1, 1),  # conv_weight
-    torch.randn(1),  # bn_weight
-    torch.randn(1),  # bn_bias
-    torch.randn(1),  # bn_running_mean
-    torch.randn(1),  # bn_running_var
-)
-
-# Example inputs for quantized and folded conv-bn2d patterns used in convert
-_quantized_conv2d_bn_example_inputs = (
-    torch.randn(1, 1, 3, 3),  # x
-    torch.randn(1, 1, 1, 1),  # conv_weight
-    torch.randn(1),  # bn_weight
-    torch.randn(1),  # bn_bias
-    torch.randn(1),  # bn_running_mean
-    torch.randn(1),  # bn_running_var
-)
-
-
 def _get_quantized_conv_bn_example_inputs_kwargs(
     is_per_channel: bool,
     has_bias: bool,
@@ -631,6 +608,28 @@ def _get_new_qspec(qspec: QuantizationSpecBase):
 
 
 def _fuse_conv_bn_qat(m: GraphModule) -> GraphModule:
+    # Example inputs for conv-bn1d patterns
+    _conv1d_bn_example_inputs = (
+        torch.randn(1, 1, 3),  # x
+        torch.randn(1, 1, 1),  # conv_weight
+        torch.randn(1),  # conv_bias
+        torch.randn(1),  # bn_weight
+        torch.randn(1),  # bn_bias
+        torch.randn(1),  # bn_running_mean
+        torch.randn(1),  # bn_running_var
+    )
+
+    # Example inputs for conv-bn2d patterns
+    _conv2d_bn_example_inputs = (
+        torch.randn(1, 1, 3, 3),  # x
+        torch.randn(1, 1, 1, 1),  # conv_weight
+        torch.randn(1),  # conv_bias
+        torch.randn(1),  # bn_weight
+        torch.randn(1),  # bn_bias
+        torch.randn(1),  # bn_running_mean
+        torch.randn(1),  # bn_running_var
+    )
+
     has_bn = any(_is_bn_node(n) for n in m.graph.nodes)
     if not has_bn:
         return m
@@ -859,6 +858,26 @@ def _copy_over_q_dq_args(original_node: Node, replacement_node: Node):
 
 
 def _fold_conv_bn_qat(m: GraphModule) -> GraphModule:
+    # Example inputs for quantized and folded conv-bn1d patterns used in convert
+    _quantized_conv1d_bn_example_inputs = (
+        torch.randn(1, 1, 3),  # x
+        torch.randn(1, 1, 1),  # conv_weight
+        torch.randn(1),  # bn_weight
+        torch.randn(1),  # bn_bias
+        torch.randn(1),  # bn_running_mean
+        torch.randn(1),  # bn_running_var
+    )
+
+    # Example inputs for quantized and folded conv-bn2d patterns used in convert
+    _quantized_conv2d_bn_example_inputs = (
+        torch.randn(1, 1, 3, 3),  # x
+        torch.randn(1, 1, 1, 1),  # conv_weight
+        torch.randn(1),  # bn_weight
+        torch.randn(1),  # bn_bias
+        torch.randn(1),  # bn_running_mean
+        torch.randn(1),  # bn_running_var
+    )
+
     has_bn = any(_is_bn_node(n) for n in m.graph.nodes)
     if not has_bn:
         return m