Repository: pytorch/functorch (language: python)

Issue:
normal with one tensor
Stack from [ghstack](https://github.com/ezyang/ghstack):
* #523
* #521
* #420
* __->__ #511




Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesRandomness.cpp b/functorch/csrc/BatchRulesRandomness.cpp
index 4fa3776ac..085e7d8d5 100644
--- a/functorch/csrc/BatchRulesRandomness.cpp
+++ b/functorch/csrc/BatchRulesRandomness.cpp
@@ -11,12 +11,22 @@
 namespace at {
 namespace functorch {
 
-void check_randomness(RandomnessType randomness) {
+void check_randomness(RandomnessType randomness, bool any_tensor_batched) {
   TORCH_CHECK(
     randomness != RandomnessType::Error,
     "vmap: called random operation while in randomness error mode. Please either use the "
     "'same' or 'different' randomness flags on vmap or perform the randomness operation out of vmap"
   );
+
+  TORCH_CHECK(
+    !(randomness == RandomnessType::Same && any_tensor_batched),
+    "Vmap does not currently support same randomness with a batched tensor input. ",
+    "Please file an issue with functorch"
+  )
+}
+
+void check_randomness(RandomnessType randomness) {
+  check_randomness(randomness, false); // for ops that don't take in any tensors, don't hit same error
 }
 
 template <typename F, F Func, typename... ExtraArgs>
@@ -81,6 +91,33 @@ Tensor randperm_batching_rule(int64_t n, ExtraArgs... extra_args) {
   }
 }
 
+template <typename F, F Func, typename... ExtraArgs>
+Tensor normal_batching_rule(const Tensor& tensor, double scalar, ExtraArgs... extra_args) {
+  c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  const auto cur_level = maybe_layer->layerId();
+
+  Tensor tensor_value;
+  optional<int64_t> tensor_bdim;
+  std::tie(tensor_value, tensor_bdim) = unwrapTensorAtLevel(tensor, cur_level);
+  tensor_value = moveBatchDimToFront(tensor_value, tensor_bdim);
+
+  RandomnessType randomness = maybe_layer->randomness();
+  check_randomness(randomness, tensor_bdim.has_value());
+  auto shape = tensor_value.sizes();
+  VmapDimVector shapeVec(shape.begin(), shape.end());
+  shapeVec.insert(shapeVec.begin(), maybe_layer->batchSize());
+
+  if (randomness == RandomnessType::Different && !tensor_bdim) {
+    tensor_value = tensor_value.unsqueeze(0).expand(shapeVec);
+  }
+  auto out = Func(tensor_value, scalar, std::forward<ExtraArgs>(extra_args)...);
+  if (randomness == RandomnessType::Same && !tensor_bdim) {
+    return out;
+  }
+  return makeBatched(out, 0, cur_level);
+}
+
 template <typename A, A a, typename C>
 struct RandomBatchRuleHelper;
 
@@ -145,6 +182,33 @@ struct RandpermBatchRuleHelper<F, Func, typelist<T1, T...>> {
   }
 };
 
+template <typename A, A a, typename C>
+struct UnaryPointwiseRandomBatchRule;
+
+template <typename F, F Func, typename A0, typename A1, typename... T>
+struct UnaryPointwiseRandomBatchRule<F, Func, typelist<A0, A1, T...>> {
+  static Tensor apply(const Tensor& tensor, double scalar, T... extra_args) {
+    return normal_batching_rule<F, Func, T...>(tensor, scalar, std::forward<T>(extra_args)...);
+  }
+};
+
+template<typename F, F Func, typename... T>
+Tensor normal_wrapper(const Tensor& tensor, double scalar, T... extra_args) {
+  return Func(scalar, tensor, extra_args...);
+}
+
+template <typename A, A a, typename C>
+struct UnaryPointwiseRandomLeadingFloatBatchRule;
+
+template <typename F, F Func, typename A0, typename A1, typename... T>
+struct UnaryPointwiseRandomLeadingFloatBatchRule<F, Func, typelist<A0, A1, T...>> {
+  static Tensor apply(double scalar, const Tensor& tensor, T... extra_args) {
+    return normal_batching_rule<decltype(&normal_wrapper<F, Func, T...>),
+                                         &normal_wrapper<F, Func, T...>,
+                                         T...>(tensor, scalar, std::forward<T>(extra_args)...);
+  }
+};
+
 TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   #define RANDOM_BATCH_RULE(op) \
     m.impl(#op, SINGLE_ARG(\
@@ -190,6 +254,16 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
       RandpermBatchRuleHelper<decltype(&ATEN_FN2(op, overload)), &ATEN_FN2(op, overload), \
                             c10::guts::function_traits<decltype(ATEN_FN2(op, overload))>::parameter_types>::apply))
 
+  #define UNARY_POINTWISE_RANDOM(op, overload) \
+    m.impl(#op"."#overload, SINGLE_ARG(\
+      UnaryPointwiseRandomBatchRule<decltype(&ATEN_FN2(op, overload)), &ATEN_FN2(op, overload), \
+                                    c10::guts::function_traits<decltype(ATEN_FN2(op, overload))>::parameter_types>::apply))
+  
+  #define UNARY_POINTWISE_RANDOM_LEADING_FLOAT(op, overload) \
+    m.impl(#op"."#overload, SINGLE_ARG(\
+      UnaryPointwiseRandomLeadingFloatBatchRule<decltype(&ATEN_FN2(op, overload)), &ATEN_FN2(op, overload), \
+                                                c10::guts::function_traits<decltype(ATEN_FN2(op, overload))>::parameter_types>::apply))
+
   RANDOM_BATCH_RULE(randn);
   RANDOM_BATCH_RULE2(randn, generator);
   RANDOM_BATCH_RULE2(randn, generator_with_names);
@@ -215,6 +289,8 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   RANDPERM_BATCH_RULE2(randperm, generator);
 
   RAND_TWO_LEADING_SCALARS_BATCH_RULE(normal, float_float);
+  UNARY_POINTWISE_RANDOM(normal, Tensor_float);
+  UNARY_POINTWISE_RANDOM_LEADING_FLOAT(normal, float_Tensor);
 
   #undef RANDOM_BATCH_RULE
   #undef RANDOM_BATCH_RULE2
@@ -225,5 +301,7 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   #undef RAND_TWO_LEADING_SCALARS_BATCH_RULE
   #undef RANDPERM_BATCH_RULE
   #undef RANDPERM_BATCH_RULE2
+  #undef UNARY_POINTWISE_RANDOM
+  #undef UNARY_POINTWISE_RANDOM_LEADING_FLOAT
 }
 }} // namespace at::functorch
diff --git a/functorch/csrc/VmapModeRegistrations.cpp b/functorch/csrc/VmapModeRegistrations.cpp
index 67fa56875..7b437f9a8 100644
--- a/functorch/csrc/VmapModeRegistrations.cpp
+++ b/functorch/csrc/VmapModeRegistrations.cpp
@@ -47,10 +47,6 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM(multinomial);
   UNSUPPORTED_RANDOM2(multinomial, out);
 
-  UNSUPPORTED_RANDOM2(normal, Tensor_float);
-  UNSUPPORTED_RANDOM2(normal, Tensor_float_out);
-  UNSUPPORTED_RANDOM2(normal, float_Tensor_out);
-  UNSUPPORTED_RANDOM2(normal, float_Tensor);
   UNSUPPORTED_RANDOM2(normal, Tensor_Tensor);
   UNSUPPORTED_RANDOM2(normal, Tensor_Tensor_out);
 


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 3d09f2524..3f3f953d2 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2603,8 +2603,6 @@ def test_no_random_op_support(self):
             (lambda t: torch.bernoulli(t, p=0.5), (torch.rand(B0, 1),)),
             (lambda t: torch.multinomial(t, 2), (torch.rand(B0, 3),)),
             (torch.normal, (torch.randn(B0, 1), torch.randn(B0, 1))),
-            (lambda t: torch.normal(t, 1.), (torch.randn(B0, 1),)),
-            (lambda t: torch.normal(0., t), (torch.randn(B0, 1),)),
             (torch.poisson, (torch.rand(B0, 1),)),
             # (torch.rand_like, (torch.rand(B0, 1),)),
             # (torch.randn_like, (torch.rand(B0, 1),)),
@@ -2616,8 +2614,6 @@ def test_no_random_op_support(self):
             (lambda t: torch.bernoulli(captured, p=0.5), (torch.rand(B0),)),
             (lambda t: torch.multinomial(captured, 2), (torch.rand(B0),)),
             (lambda t: torch.normal(captured, captured), (torch.randn(B0),)),
-            (lambda t: torch.normal(captured, 1.), (torch.randn(B0),)),
-            (lambda t: torch.normal(0., captured), (torch.randn(B0),)),
             (lambda t: torch.poisson(captured), (torch.rand(B0),)),
             # (lambda t: torch.rand_like(captured), (torch.rand(B0),)),
             # (lambda t: torch.randn_like(captured) , (torch.rand(B0),)),
@@ -3413,12 +3409,12 @@ def f(x, y):
         y = torch.randn(2, 3, device=device)
         self.assertTrue(isinstance(vmap(f)(x, y), Point))
 
+    def reset_random(self, generator, orig_state, use_generator, seed):
+        return generator.set_state(orig_state) if use_generator else torch.manual_seed(seed)
+
     @parametrize('randomness', ['same', 'different', 'error'])
     @parametrize('use_generator', [True, False])
     def test_random_behavior(self, device, randomness, use_generator):
-        def reset_random():
-            return generator.set_state(orig_state) if use_generator else torch.manual_seed(seed)
-
         generator = torch.Generator(device=device)
         orig_state = generator.get_state()
         kwargs = {'device': device, 'generator': generator} if use_generator else {'device': device}
@@ -3443,16 +3439,16 @@ def reset_random():
                 return
 
             passed = torch.randn(B0, B0, device=device)
-            generator = reset_random()
+            generator = self.reset_random(generator, orig_state, use_generator, seed)
             vmap_result = vmap(op, in_dims=(0, None), randomness=randomness)(passed, [B0])
             if randomness == "different":
                 passed = torch.randn([B0, B0], device=device)  # reset for in place operation
-                generator = reset_random()
+                generator = self.reset_random(generator, orig_state, use_generator, seed)
                 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()
+                generator = self.reset_random(generator, orig_state, use_generator, seed)
                 expected = op(passed, [B0])
                 for i in range(B0):
                     assert torch.allclose(vmap_result[i], expected)
@@ -3527,6 +3523,49 @@ def test_random_inplace_not_batched(self, device, use_generator, randomness):
             op(passed, vmaped_value)
             assert torch.allclose(unvmaped_value, passed)
 
+    @parametrize('randomness', ['same', 'different', 'error'])
+    @parametrize('use_generator', [True, False])
+    @parametrize('batched_input', [True, False])
+    def test_distributions(self, device, randomness, use_generator, batched_input):
+        supported_ops = [
+            lambda _, t: torch.normal(t, 1., **kwargs),
+            lambda _, t: torch.normal(0., torch.abs(t), **kwargs),
+        ]
+
+        B0 = 4
+        seed = 1234567
+        generator = torch.Generator(device=device)
+        orig_state = generator.get_state()
+        generator = self.reset_random(generator, orig_state, use_generator, seed)
+
+        in_dims = (0, 0) if batched_input else (0, None)
+        args_shape = ((B0, B0), (B0, B0)) if batched_input else ((B0, B0), (B0,))
+        kwargs = {'generator': generator} if use_generator else {}
+
+        for op in supported_ops:
+            args = tuple(torch.randn(*shape, device=device) for shape in args_shape)
+            if randomness == 'error' or (randomness == 'same' and batched_input):
+                error_regex = r"called random operation while in randomness error mode"
+                same_random_regex = r"Vmap does not currently support same randomness with a batched tensor input"
+                regex = error_regex if randomness == 'error' else same_random_regex
+
+                with self.assertRaisesRegex(RuntimeError, regex):
+                    vmap(op, in_dims=in_dims, randomness=randomness)(*args)
+                return
+
+            generator = self.reset_random(generator, orig_state, use_generator, seed)
+            vmap_result = vmap(op, in_dims=in_dims, randomness=randomness)(*args)
+            if randomness == "different":
+                generator = self.reset_random(generator, orig_state, use_generator, seed)
+                args = args if batched_input else (args[0], args[1].unsqueeze(0).expand(B0, B0))
+                expected = op(*args)
+                assert torch.allclose(vmap_result, expected)
+            else:
+                generator = self.reset_random(generator, orig_state, use_generator, seed)
+                expected = op(*args)
+                for i in range(B0):
+                    assert torch.allclose(vmap_result[i], expected)
+
 
 only_for = ("cpu", "cuda")
 instantiate_device_type_tests(TestVmapOperatorsOpInfo, globals(), only_for=only_for)