Revert "Break graph on manual_seed. (pytorch#107594)"

This reverts commit 6ad5568.

Reverted pytorch#107594 on behalf of https://github.com/huydhn due to Sorry for reverting your change, but it has an import issue that breaks internal code ([comment](pytorch#107594 (comment)))

Author: pytorchmergebot

test/dynamo/test_functions.py

@@ -1146,15 +1146,6 @@ def test_numpy_random():
         x = np.random.randn(2, 2)
         return x - x
 
-    def test_manual_seed(self):
-        @torch.compile
-        def foo():
-            torch.manual_seed(3)
-            return torch.randint(0, 5, (5,))
-
-        self.assertEqual(foo(), foo())
-        self.assertEqual(foo(), foo())
-
 
 def global_func_with_default_tensor_args(
     x=torch.zeros((2, 2)), *, kw_x=torch.zeros((1, 2))

test/dynamo/test_misc.py

@@ -1813,6 +1813,14 @@ def fn(x, obj):
         res = opt_fn(x, obj)
         self.assertTrue(same(ref, res))
 
+    def test_manual_seed(self):
+        def fn(a, b):
+            x = a + b
+            torch.manual_seed(9000)
+            return x + 1
+
+        torch._dynamo.testing.standard_test(self, fn=fn, nargs=2, expected_ops=3)
+
     def test_usr_cls_staticmethod(self):
         class Foo:
             @staticmethod
@@ -2259,17 +2267,13 @@ def fn(x):
             torch.manual_seed(attention_seed)
             return (x,)
 
-        x = torch.randn(10, requires_grad=True)
+        x = torch.randn(100, requires_grad=True)
         ref = fn(x)
 
-        # Python code is needed here, since torch.manual_seed graph-breaks.
-        # Refs: https://github.com/pytorch/pytorch/issues/107187
-        opt_fn = torch._dynamo.optimize(cnts, nopython=False)(fn)
+        opt_fn = torch._dynamo.optimize(cnts, nopython=True)(fn)
         res = opt_fn(x)
 
         self.assertTrue(same(ref, res))
-        self.assertEqual(cnts.op_count, 1)
-        self.assertEqual(cnts.frame_count, 1)
 
     def test_is_tensor_like(self):
         cnts = torch._dynamo.testing.CompileCounter()

test/inductor/test_torchinductor_opinfo.py

@@ -262,6 +262,7 @@ def format_op(op):
     "nn.functional.instance_norm": {f16},
     "nn.functional.local_response_norm": {f16},
     "nn.functional.normalize": {f16},
+    "nn.functional.rrelu": {f16, f32, f64},
     "nn.functional.soft_margin_loss": {f16},
     "nn.functional.softsign": {f16},
     "nn.functional.triplet_margin_loss": {f16},
@@ -277,6 +278,7 @@ def format_op(op):
     "sparse.sampled_addmm": {f32, f64},
     ("std_mean", "unbiased"): {f16},
     "to_sparse": {f16, f32, f64},
+    "uniform": {f16, f32, f64},
 }
 
 
@@ -340,13 +342,15 @@ def get_skips_and_xfails(from_dict, xfails=True):
 )
 
 
-def wrapper_noop_set_seed(op, *args, **kwargs):
+def wrapper_set_seed(op, *args, **kwargs):
+    """Wrapper to set seed manually for some functions like dropout
+    See: https://github.com/pytorch/pytorch/pull/62315#issuecomment-896143189 for more details.
+    """
+    torch.manual_seed(42)
     return op(*args, **kwargs)
 
 
-torch.testing._internal.common_methods_invocations.wrapper_set_seed = (
-    wrapper_noop_set_seed
-)
+torch.testing._internal.common_methods_invocations.wrapper_set_seed = wrapper_set_seed
 
 # This file does a global patch to `disable_global_flags()` - which we should not invoke in non testing cases.
 torch._dynamo.variables.torch.tensor_dunder_fns.append(

torch/__init__.py

@@ -1303,17 +1303,6 @@ def _dtype(self):
 # The _tensor_classes set is initialized by the call to _C._initialize_tensor_type_bindings()
 _tensor_classes: Set[Type] = set()
 
-################################################################################
-# Import TorchDynamo's lazy APIs to avoid circular dependenices
-################################################################################
-
-# needs to be before from .functional import * to avoid circular dependencies
-from ._compile import _disable_dynamo
-
-################################################################################
-# Import miscelaneous torch functions
-################################################################################
-
 # If you edit these imports, please update torch/__init__.py.in as well
 from .random import set_rng_state, get_rng_state, manual_seed, initial_seed, seed
 from .serialization import save, load
@@ -1378,6 +1367,13 @@ def manager_path():
 
 
 
+################################################################################
+# Import TorchDynamo's lazy APIs to avoid circular dependenices
+################################################################################
+
+# needs to be before from .functional import * to avoid circular dependencies
+from ._compile import _disable_dynamo
+
 ################################################################################
 # Import interface functions defined in Python
 ################################################################################

torch/_dynamo/variables/torch.py

@@ -453,9 +453,6 @@ def call_function(
         elif self.value is torch.nn.Parameter:
             # https://github.com/pytorch/pytorch/issues/99569
             unimplemented("torch.nn.Parameter not supported")
-        elif self.value is torch.manual_seed:
-            # https://github.com/pytorch/pytorch/issues/107187
-            unimplemented("torch.manual_seed not supported")
         if (
             self.value.__name__ == "get_state"
             and hasattr(self.value, "__self__")

torch/random.py

@@ -23,7 +23,6 @@ def get_rng_state() -> torch.Tensor:
     return default_generator.get_state()
 
 
-@torch._disable_dynamo
 def manual_seed(seed) -> torch._C.Generator:
     r"""Sets the seed for generating random numbers. Returns a
     `torch.Generator` object.