Repository: pytorch/functorch (language: python)

Issue:
random_ and normal_ support
Stack from [ghstack](https://github.com/ezyang/ghstack):
* #523
* #521
* #420
* #511
* #510
* __->__ #412




Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesRandomness.cpp b/functorch/csrc/BatchRulesRandomness.cpp
index 3763a612d..6047e1467 100644
--- a/functorch/csrc/BatchRulesRandomness.cpp
+++ b/functorch/csrc/BatchRulesRandomness.cpp
@@ -35,6 +35,32 @@ Tensor random_batching_rule(IntArrayRef shape, ExtraArgs... extra_args) {
   }
 }
 
+template <typename F, F Func, typename... ExtraArgs>
+Tensor& random_inplace_batching_rule(Tensor& self, ExtraArgs... extra_args) {
+  c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  const auto cur_level = maybe_layer->layerId();
+  Tensor self_value;
+  optional<int64_t> self_bdim;
+  std::tie(self_value, self_bdim) = unwrapTensorAtLevel(self, cur_level);
+  self_value = moveBatchDimToFront(self_value, self_bdim);
+  RandomnessType randomness = maybe_layer->randomness();
+  check_randomness(randomness);
+  TORCH_CHECK(
+    !(randomness == RandomnessType::Different && !self_bdim),
+    "vmap: Cannot ask for different inplace randomness on an unbatched tensor. This will appear like same randomness. ",
+    "If this is necessary for your usage, please file an issue with functorch.");
+  if (randomness == RandomnessType::Same && self_bdim) {
+    auto intermediate = empty(self.sizes(), self.options());
+    Func(intermediate, std::forward<ExtraArgs>(extra_args)...);
+    self.copy_(intermediate); // batching should make this just work out...
+    return self;
+  } else {
+    Func(self_value, std::forward<ExtraArgs>(extra_args)...);
+    return self;
+  }
+}
+
 template <typename F, F Func, typename... ExtraArgs>
 Tensor randperm_batching_rule(int64_t n, ExtraArgs... extra_args) {
   c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
@@ -70,6 +96,16 @@ Tensor rand_int_wrapper(IntArrayRef shape, int64_t high, T... extra_args) {
   return Func(high, shape, std::forward<T>(extra_args)...);
 }
 
+template <typename A, A a, typename C>
+struct RandomInplaceBatchRuleHelper;
+
+template <typename F, F Func, typename T1, typename... T>
+struct RandomInplaceBatchRuleHelper<F, Func, typelist<T1, T...>> {
+  static Tensor& apply(Tensor& self, T... extra_args) {
+    return random_inplace_batching_rule<F, Func, T...>(self, std::forward<T>(extra_args)...);
+  }
+};
+
 template <typename A, A a, typename C>
 struct RandIntBatchRuleHelper;
 
@@ -119,6 +155,16 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
     m.impl(#op"."#overload, SINGLE_ARG(\
       RandomBatchRuleHelper<decltype(&ATEN_FN2(op, overload)), &ATEN_FN2(op, overload), \
                             c10::guts::function_traits<decltype(ATEN_FN2(op, overload))>::parameter_types>::apply))
+  
+  #define RANDOM_INPLACE_BATCH_RULE(op) \
+    m.impl(#op, SINGLE_ARG(\
+      RandomInplaceBatchRuleHelper<decltype(&ATEN_FN(op)), &ATEN_FN(op), \
+                            c10::guts::function_traits<decltype(ATEN_FN(op))>::parameter_types>::apply))
+
+  #define RANDOM_INPLACE_BATCH_RULE2(op, overload) \
+    m.impl(#op"."#overload, SINGLE_ARG(\
+      RandomInplaceBatchRuleHelper<decltype(&ATEN_FN2(op, overload)), &ATEN_FN2(op, overload), \
+                            c10::guts::function_traits<decltype(ATEN_FN2(op, overload))>::parameter_types>::apply))
 
   #define RANDINT_BATCH_RULE(op) \
     m.impl(#op, SINGLE_ARG(\
@@ -154,6 +200,12 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   RANDOM_BATCH_RULE2(rand, generator_with_names);
   RANDOM_BATCH_RULE2(rand, names);
 
+  RANDOM_INPLACE_BATCH_RULE(random_);
+  RANDOM_INPLACE_BATCH_RULE2(random_, from);
+  RANDOM_INPLACE_BATCH_RULE2(random_, to);
+
+  RANDOM_INPLACE_BATCH_RULE(normal_);
+
   RANDINT_BATCH_RULE(randint);
   RANDINT_BATCH_RULE2(randint, generator);
   RANDINT_LOW_BATCH_RULE(randint, low);
@@ -164,6 +216,8 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
 
   #undef RANDOM_BATCH_RULE
   #undef RANDOM_BATCH_RULE2
+  #undef RANDOM_INPLACE_BATCH_RULE
+  #undef RANDOM_INPLACE_BATCH_RULE2
   #undef RANDINT_BATCH_RULE
   #undef RANDINT_BATCH_RULE2
   #undef RANDINT_LOW_BATCH_RULE
diff --git a/functorch/csrc/VmapModeRegistrations.cpp b/functorch/csrc/VmapModeRegistrations.cpp
index ce3f4d01c..5a3d483d8 100644
--- a/functorch/csrc/VmapModeRegistrations.cpp
+++ b/functorch/csrc/VmapModeRegistrations.cpp
@@ -55,14 +55,9 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM2(normal, Tensor_Tensor_out);
   UNSUPPORTED_RANDOM2(normal, float_float);
   UNSUPPORTED_RANDOM2(normal, float_float_out);
-  UNSUPPORTED_RANDOM(normal_);
 
   UNSUPPORTED_RANDOM(poisson);
 
-  UNSUPPORTED_RANDOM2(random_, from);
-  UNSUPPORTED_RANDOM2(random_, to);
-  UNSUPPORTED_RANDOM(random_);
-
   UNSUPPORTED_RANDOM(randint_like);
   UNSUPPORTED_RANDOM2(randint_like, low_dtype);
 


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index d094a71d8..7ae0ed8dc 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2630,10 +2630,6 @@ def test_no_random_op_support(self):
             (lambda t: t.exponential_(), (torch.randn(B0, 1),)),
             (lambda t: t.geometric_(0.5), (torch.randn(B0, 1),)),
             (lambda t: t.log_normal_(), (torch.randn(B0, 1),)),
-            (lambda t: t.normal_(), (torch.randn(B0, 1),)),
-            (lambda t: t.random_(), (torch.randn(B0, 1),)),
-            (lambda t: t.random_(0, 2), (torch.randn(B0, 1),)),
-            (lambda t: t.random_(2), (torch.randn(B0, 1),)),
             (lambda t: t.uniform_(), (torch.randn(B0, 1),)),
 
             # in-place on captured tensor
@@ -2642,10 +2638,6 @@ def test_no_random_op_support(self):
             (lambda t: captured.exponential_(), (torch.randn(B0),)),
             (lambda t: captured.geometric_(0.5), (torch.randn(B0),)),
             (lambda t: captured.log_normal_(), (torch.randn(B0),)),
-            (lambda t: captured.normal_(), (torch.randn(B0),)),
-            (lambda t: captured.random_(), (torch.randn(B0),)),
-            (lambda t: captured.random_(0, 2), (torch.randn(B0),)),
-            (lambda t: captured.random_(2), (torch.randn(B0),)),
             (lambda t: captured.uniform_(), (torch.randn(B0),)),
         ]
         for op, args in random_ops:
@@ -3430,11 +3422,13 @@ def reset_random():
         generator = torch.Generator(device=device)
         orig_state = generator.get_state()
         kwargs = {'device': device, 'generator': generator} if use_generator else {'device': device}
+        only_gen_kwarg = {'generator': generator} if use_generator else {}
         supported_random_ops = [
             lambda _, shape: torch.randn(shape, **kwargs),
             lambda _, shape: torch.rand(shape, **kwargs),
             lambda _, shape: torch.randint(100, shape, **kwargs),
             lambda _, shape: torch.randint(5, 100, shape, **kwargs),
+            lambda t, _: t.random_(**only_gen_kwarg),
         ]
 
         B0 = 4
@@ -3447,14 +3441,17 @@ def reset_random():
                     vmap(op, in_dims=(0, None), randomness=randomness)(passed, [B0])
                 return
 
+            passed = torch.randn(B0, B0, device=device)
             generator = reset_random()
             vmap_result = vmap(op, in_dims=(0, None), randomness=randomness)(passed, [B0])
-
-            generator = reset_random()
-            if randomness == 'different':
+            if randomness == "different":
+                passed = torch.randn([B0, B0], device=device)  # reset for in place operation
+                generator = reset_random()
                 expected = op(passed, [B0, B0])
                 assert torch.allclose(vmap_result, expected)
             else:
+                passed = torch.randn(B0, device=device)  # reset for in place operation
+                generator = reset_random()
                 expected = op(passed, [B0])
                 for i in range(B0):
                     assert torch.allclose(vmap_result[i], expected)
@@ -3490,6 +3487,45 @@ def test_randperm(self, device, randomness, use_generator):
             for i in range(B0):
                 assert torch.allclose(vmap_result[i], expected)
 
+    @parametrize('use_generator', [True, False])
+    @parametrize('randomness', ['error', 'same', 'different'])
+    def test_random_inplace_not_batched(self, device, use_generator, randomness):
+        # tests that when in place random ops are being called during vmap but not with a batched tensor,
+        # it behaves like a normal random_ when using same random and errors during different random
+        supported_ops = [
+            lambda t, _: t.random_(**kwargs),
+            lambda t, _: t.random_(100, **kwargs),
+            lambda t, _: t.random_(-5, 100, **kwargs),
+            lambda t, _: t.normal_(**kwargs),
+        ]
+
+        B0 = 4
+        seed = 1234567
+        generator = torch.Generator(device=device)
+        orig_state = generator.get_state()
+        kwargs = {'generator': generator} if use_generator else {}
+
+        for op in supported_ops:
+            vmaped_value, unvmaped_value = torch.randn(B0, B0, device=device), torch.randn(B0, B0, device=device)
+            generator = generator.set_state(orig_state)
+            torch.manual_seed(seed)
+            if randomness == 'error' or randomness == 'different':
+                error_regex = r"called random operation while in randomness error mode"
+                different_random_regex = r"different inplace randomness on an unbatched tensor"
+                regex = error_regex if randomness == 'error' else different_random_regex
+
+                with self.assertRaisesRegex(RuntimeError, regex):
+                    vmap(op, in_dims=(None, 0), randomness=randomness)(unvmaped_value, vmaped_value)
+                return
+
+            vmap(op, in_dims=(None, 0), randomness=randomness)(unvmaped_value, vmaped_value)
+
+            passed = torch.randn([B0, B0], device=device)  # reset for in place operation
+            generator = generator.set_state(orig_state)
+            torch.manual_seed(seed)
+            op(passed, vmaped_value)
+            assert torch.allclose(unvmaped_value, passed)
+
 
 only_for = ("cpu", "cuda")
 instantiate_device_type_tests(TestVmapOperatorsOpInfo, globals(), only_for=only_for)