Repository: pytorch/functorch (language: python)

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




Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesRandomness.cpp b/functorch/csrc/BatchRulesRandomness.cpp
index 6047e1467..4fa3776ac 100644
--- a/functorch/csrc/BatchRulesRandomness.cpp
+++ b/functorch/csrc/BatchRulesRandomness.cpp
@@ -118,20 +118,20 @@ struct RandIntBatchRuleHelper<F, Func, typelist<T1, T2, T...>> {
   }
 };
 
-template <typename F, F Func, typename... T>
-Tensor rand_int_low_wrapper(IntArrayRef shape, int64_t high, int64_t low, T... extra_args) {
-  return Func(high, low, shape, std::forward<T>(extra_args)...);
+template <typename F, F Func, typename T0, typename T1, typename... T>
+Tensor rand_int_low_wrapper(IntArrayRef shape, T0 scalar0, T1 scalar1, T... extra_args) {
+  return Func(scalar0, scalar1, shape, std::forward<T>(extra_args)...);
 }
 
 template <typename A, A a, typename C>
-struct RandIntLowBatchRuleHelper;
-
-template <typename F, F Func, typename T1, typename T2, typename T3, typename... T>
-struct RandIntLowBatchRuleHelper<F, Func, typelist<T1, T2, T3, T...>> {
-  static Tensor apply(int64_t high, int64_t low, IntArrayRef shape, T... extra_args) {
-    return random_batching_rule<decltype(&rand_int_low_wrapper<F, Func, T...>),
-                                &rand_int_low_wrapper<F, Func, T...>,
-                                int64_t, int64_t, T...>(shape, high, low, std::forward<T>(extra_args)...);
+struct RandTwoLeadingScalarsBatchRuleHelper;
+
+template <typename F, F Func, typename T0, typename T1, typename T2, typename... T>
+struct RandTwoLeadingScalarsBatchRuleHelper<F, Func, typelist<T0, T1, T2, T...>> {
+  static Tensor apply(T0 scalar0, T1 scalar1, IntArrayRef shape, T... extra_args) {
+    return random_batching_rule<decltype(&rand_int_low_wrapper<F, Func, T0, T1, T...>),
+                                &rand_int_low_wrapper<F, Func, T0, T1, T...>,
+                                int64_t, int64_t, T...>(shape, scalar0, scalar1, std::forward<T>(extra_args)...);
   }
 };
 
@@ -176,9 +176,9 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
       RandIntBatchRuleHelper<decltype(&ATEN_FN2(op, overload)), &ATEN_FN2(op, overload), \
                             c10::guts::function_traits<decltype(ATEN_FN2(op, overload))>::parameter_types>::apply))
 
-  #define RANDINT_LOW_BATCH_RULE(op, overload) \
+  #define RAND_TWO_LEADING_SCALARS_BATCH_RULE(op, overload) \
     m.impl(#op"."#overload, SINGLE_ARG(\
-      RandIntLowBatchRuleHelper<decltype(&ATEN_FN2(op, overload)), &ATEN_FN2(op, overload), \
+      RandTwoLeadingScalarsBatchRuleHelper<decltype(&ATEN_FN2(op, overload)), &ATEN_FN2(op, overload), \
                                 c10::guts::function_traits<decltype(ATEN_FN2(op, overload))>::parameter_types>::apply))
   #define RANDPERM_BATCH_RULE(op) \
     m.impl(#op, SINGLE_ARG(\
@@ -208,19 +208,21 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
 
   RANDINT_BATCH_RULE(randint);
   RANDINT_BATCH_RULE2(randint, generator);
-  RANDINT_LOW_BATCH_RULE(randint, low);
-  RANDINT_LOW_BATCH_RULE(randint, low_generator);
+  RAND_TWO_LEADING_SCALARS_BATCH_RULE(randint, low);
+  RAND_TWO_LEADING_SCALARS_BATCH_RULE(randint, low_generator);
 
   RANDPERM_BATCH_RULE(randperm);
   RANDPERM_BATCH_RULE2(randperm, generator);
 
+  RAND_TWO_LEADING_SCALARS_BATCH_RULE(normal, float_float);
+
   #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
+  #undef RAND_TWO_LEADING_SCALARS_BATCH_RULE
   #undef RANDPERM_BATCH_RULE
   #undef RANDPERM_BATCH_RULE2
 }
diff --git a/functorch/csrc/VmapModeRegistrations.cpp b/functorch/csrc/VmapModeRegistrations.cpp
index 5a3d483d8..67fa56875 100644
--- a/functorch/csrc/VmapModeRegistrations.cpp
+++ b/functorch/csrc/VmapModeRegistrations.cpp
@@ -53,8 +53,6 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM2(normal, float_Tensor);
   UNSUPPORTED_RANDOM2(normal, Tensor_Tensor);
   UNSUPPORTED_RANDOM2(normal, Tensor_Tensor_out);
-  UNSUPPORTED_RANDOM2(normal, float_float);
-  UNSUPPORTED_RANDOM2(normal, float_float_out);
 
   UNSUPPORTED_RANDOM(poisson);
 


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 7ae0ed8dc..3d09f2524 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -3429,6 +3429,7 @@ def reset_random():
             lambda _, shape: torch.randint(100, shape, **kwargs),
             lambda _, shape: torch.randint(5, 100, shape, **kwargs),
             lambda t, _: t.random_(**only_gen_kwarg),
+            lambda _, shape: torch.normal(0., 1., shape, **kwargs),
         ]
 
         B0 = 4