Repository: pytorch/functorch (language: python)

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




Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesBinaryOps.cpp b/functorch/csrc/BatchRulesBinaryOps.cpp
index e343ad897..7b014f9f1 100644
--- a/functorch/csrc/BatchRulesBinaryOps.cpp
+++ b/functorch/csrc/BatchRulesBinaryOps.cpp
@@ -89,6 +89,48 @@ struct BinaryPointwiseBatchRuleHelper<F, Func, typelist<T1, T2, T...>> {
       &fn,\
       c10::guts::function_traits<decltype(fn)>::parameter_types>::apply)
 
+template <typename A, A a, typename C>
+struct BinaryRandomPointwiseBatchRuleHelper;
+
+template <typename F, F Func, typename T1, typename T2, typename... T>
+struct BinaryRandomPointwiseBatchRuleHelper<F, Func, typelist<T1, T2, T...>> {
+  static Tensor apply(const Tensor& tensor, const Tensor& other, T... extra_args) {
+    c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
+    auto maybe_layer = maybeCurrentDynamicLayer();
+    auto cur_level = maybe_layer->layerId();
+    RandomnessType randomness = maybe_layer->randomness();
+
+    Tensor tensor_value;
+    optional<int64_t> tensor_bdim;
+    std::tie(tensor_value, tensor_bdim) = unwrapTensorAtLevel(tensor, cur_level);
+
+    Tensor other_value;
+    optional<int64_t> other_bdim;
+    std::tie(other_value, other_bdim) = unwrapTensorAtLevel(other, cur_level);
+
+    check_randomness(randomness, (tensor_bdim || other_bdim));
+    if (randomness == RandomnessType::Different && !tensor_bdim && !other_bdim) {
+      auto shape = tensor_value.sizes();
+      VmapDimVector shapeVec(shape.begin(), shape.end());
+      shapeVec.insert(shapeVec.begin(), maybe_layer->batchSize());
+      tensor_value = tensor_value.unsqueeze(0).expand(shapeVec);
+      tensor_bdim = 0;
+    } else if (randomness == RandomnessType::Same && !tensor_bdim && !other_bdim) {
+      return Func(tensor_value, other_value, std::forward<T>(extra_args)...);
+    }
+    auto res = _binary_pointwise_batch_rule<F, Func, T...>(
+      tensor_value, tensor_bdim, other_value, other_bdim,
+      std::forward<T>(extra_args)...);
+    return makeBatched(std::get<0>(res), std::get<1>(res), cur_level);
+  }
+};
+
+#define BINARY_RANDOM_POINTWISE_BATCH_RULE(fn) SINGLE_ARG(\
+    BinaryRandomPointwiseBatchRuleHelper<\
+      decltype(&fn),\
+      &fn,\
+      c10::guts::function_traits<decltype(fn)>::parameter_types>::apply)
+
 template <typename M, M Meth, typename... ExtraArgs>
 void binary_pointwise_inplace_batch_rule(
     Tensor& tensor, optional<int64_t> tensor_batch_dim,
@@ -220,6 +262,12 @@ std::tuple<Tensor,optional<int64_t>> cdist_backward_batch_rule(
   return std::make_tuple(out, out_bdim);
 }
 
+TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
+#define BINARY_RANDOM_POINTWISE2(op, overload) \
+  m.impl(#op"."#overload, BINARY_RANDOM_POINTWISE_BATCH_RULE(ATEN_FN2(op, overload)));
+
+  BINARY_RANDOM_POINTWISE2(normal, Tensor_Tensor);
+}
 
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
 #define BINARY_POINTWISE2(op, overload) \
diff --git a/functorch/csrc/BatchRulesHelper.cpp b/functorch/csrc/BatchRulesHelper.cpp
index a0e03fc60..81a39d243 100644
--- a/functorch/csrc/BatchRulesHelper.cpp
+++ b/functorch/csrc/BatchRulesHelper.cpp
@@ -83,6 +83,24 @@ Tensor maybePadToLogicalRank(const Tensor& tensor, optional<int64_t> has_bdim, i
   return tensor.view(new_sizes);
 }
 
+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
+}
+
 Tensor reshape_dim_into(int64_t src, int64_t dst, const Tensor& x) {
   auto x_dim = x.dim();
   src = maybe_wrap_dim(src, x_dim);
diff --git a/functorch/csrc/BatchRulesHelper.h b/functorch/csrc/BatchRulesHelper.h
index ad159a778..0cb2f151a 100644
--- a/functorch/csrc/BatchRulesHelper.h
+++ b/functorch/csrc/BatchRulesHelper.h
@@ -34,6 +34,9 @@ void vmapIncompatibleInplaceError(const char* schema_name);
 
 Tensor maybePadToLogicalRank(const Tensor& tensor, optional<int64_t> has_bdim, int64_t logical_rank);
 
+void check_randomness(RandomnessType randomness);
+void check_randomness(RandomnessType randomness, bool any_tensor_bdim);
+
 inline Tensor ensure_has_bdim(const Tensor& tensor, bool has_bdim, int64_t batch_size) {
   if (has_bdim) {
     return tensor;
diff --git a/functorch/csrc/BatchRulesRandomness.cpp b/functorch/csrc/BatchRulesRandomness.cpp
index 085e7d8d5..6da49b276 100644
--- a/functorch/csrc/BatchRulesRandomness.cpp
+++ b/functorch/csrc/BatchRulesRandomness.cpp
@@ -11,24 +11,6 @@
 namespace at {
 namespace functorch {
 
-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>
 Tensor random_batching_rule(IntArrayRef shape, ExtraArgs... extra_args) {
   c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
diff --git a/functorch/csrc/VmapModeRegistrations.cpp b/functorch/csrc/VmapModeRegistrations.cpp
index 7b437f9a8..da3790b92 100644
--- a/functorch/csrc/VmapModeRegistrations.cpp
+++ b/functorch/csrc/VmapModeRegistrations.cpp
@@ -47,9 +47,6 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM(multinomial);
   UNSUPPORTED_RANDOM2(multinomial, out);
 
-  UNSUPPORTED_RANDOM2(normal, Tensor_Tensor);
-  UNSUPPORTED_RANDOM2(normal, Tensor_Tensor_out);
-
   UNSUPPORTED_RANDOM(poisson);
 
   UNSUPPORTED_RANDOM(randint_like);


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 3f3f953d2..f94a142a4 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2602,7 +2602,6 @@ def test_no_random_op_support(self):
             (torch.bernoulli, (torch.rand(B0, 1),)),
             (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))),
             (torch.poisson, (torch.rand(B0, 1),)),
             # (torch.rand_like, (torch.rand(B0, 1),)),
             # (torch.randn_like, (torch.rand(B0, 1),)),
@@ -2613,7 +2612,6 @@ def test_no_random_op_support(self):
             (lambda t: torch.bernoulli(captured), (torch.rand(B0),)),
             (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.poisson(captured), (torch.rand(B0),)),
             # (lambda t: torch.rand_like(captured), (torch.rand(B0),)),
             # (lambda t: torch.randn_like(captured) , (torch.rand(B0),)),
@@ -3426,6 +3424,7 @@ def test_random_behavior(self, device, randomness, use_generator):
             lambda _, shape: torch.randint(5, 100, shape, **kwargs),
             lambda t, _: t.random_(**only_gen_kwarg),
             lambda _, shape: torch.normal(0., 1., shape, **kwargs),
+            lambda t, _: t.normal_(**only_gen_kwarg),
         ]
 
         B0 = 4
@@ -3530,6 +3529,7 @@ 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),
+            lambda _, t: torch.normal(t, torch.abs(t), **kwargs),
         ]
 
         B0 = 4