Repository: pytorch/functorch (language: python)

Issue:
randint support
Stack from [ghstack](https://github.com/ezyang/ghstack):
* #420
* #412
* #411
* __->__ #410
* #409




Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesRandomness.cpp b/functorch/csrc/BatchRulesRandomness.cpp
index 02e27c34e..9e2cc4055 100644
--- a/functorch/csrc/BatchRulesRandomness.cpp
+++ b/functorch/csrc/BatchRulesRandomness.cpp
@@ -30,8 +30,7 @@ Tensor random_batching_rule(IntArrayRef shape, ExtraArgs... extra_args) {
   if (randomness == "different") {
     return makeBatched(Func(shapeVec, std::forward<ExtraArgs>(extra_args)...), 0, maybe_layer->layerId());
   } else {
-    const auto res = Func(shape, std::forward<ExtraArgs>(extra_args)...);
-    return makeBatched(res.unsqueeze(0).expand(shapeVec), 0, maybe_layer->layerId());
+    return Func(shape, std::forward<ExtraArgs>(extra_args)...);
   }
 }
 
@@ -45,6 +44,39 @@ struct RandomBatchRuleHelper<F, Func, typelist<T1, T...>> {
   }
 };
 
+template <typename F, F Func, typename... T>
+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 RandIntBatchRuleHelper;
+
+template <typename F, F Func, typename T1, typename T2, typename... T>
+struct RandIntBatchRuleHelper<F, Func, typelist<T1, T2, T...>> {
+  static Tensor apply(int64_t high, IntArrayRef shape, T... extra_args) {
+    return random_batching_rule<decltype(&rand_int_wrapper<F, Func, T...>),
+                                &rand_int_wrapper<F, Func, T...>,
+                                int64_t, T...>(shape, high, std::forward<T>(extra_args)...);
+  }
+};
+
+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 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)...);
+  }
+};
 
 TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   #define RANDOM_BATCH_RULE(op) \
@@ -57,6 +89,21 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
       RandomBatchRuleHelper<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(\
+      RandIntBatchRuleHelper<decltype(&ATEN_FN(op)), &ATEN_FN(op), \
+                             c10::guts::function_traits<decltype(ATEN_FN(op))>::parameter_types>::apply))
+
+  #define RANDINT_BATCH_RULE2(op, overload) \
+    m.impl(#op"."#overload, SINGLE_ARG(\
+      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) \
+    m.impl(#op"."#overload, SINGLE_ARG(\
+      RandIntLowBatchRuleHelper<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);
@@ -67,7 +114,15 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   RANDOM_BATCH_RULE2(rand, generator_with_names);
   RANDOM_BATCH_RULE2(rand, names);
 
+  RANDINT_BATCH_RULE(randint);
+  RANDINT_BATCH_RULE2(randint, generator);
+  RANDINT_LOW_BATCH_RULE(randint, low);
+  RANDINT_LOW_BATCH_RULE(randint, low_generator);
+
   #undef RANDOM_BATCH_RULE
   #undef RANDOM_BATCH_RULE2
+  #undef RANDINT_BATCH_RULE
+  #undef RANDINT_BATCH_RULE2
+  #undef RANDINT_LOW_BATCH_RULE
 }
 }} // namespace at::functorch
\ No newline at end of file
diff --git a/functorch/csrc/VmapModeRegistrations.cpp b/functorch/csrc/VmapModeRegistrations.cpp
index d2c0a8cc4..893697675 100644
--- a/functorch/csrc/VmapModeRegistrations.cpp
+++ b/functorch/csrc/VmapModeRegistrations.cpp
@@ -71,15 +71,6 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM2(randperm, out);
   UNSUPPORTED_RANDOM2(randperm, generator_out);
 
-  UNSUPPORTED_RANDOM(randint);
-  UNSUPPORTED_RANDOM2(randint, generator);
-  UNSUPPORTED_RANDOM2(randint, low);
-  UNSUPPORTED_RANDOM2(randint, low_generator);
-  UNSUPPORTED_RANDOM2(randint, out);
-  UNSUPPORTED_RANDOM2(randint, generator_out);
-  UNSUPPORTED_RANDOM2(randint, low_out);
-  UNSUPPORTED_RANDOM2(randint, low_generator_out);
-
   UNSUPPORTED_RANDOM(uniform_);
 }
 


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 279efb8e5..3abb0f05b 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2649,7 +2649,6 @@ def test_no_random_op_support(self):
             (lambda t: captured.uniform_(), (torch.randn(B0),)),
 
             # factory functions
-            (lambda t: torch.randint(5, [1]), (torch.randn(B0),)),
             (lambda t: torch.randperm(5), (torch.randn(B0),)),
         ]
         for op, args in random_ops:
@@ -3438,6 +3437,8 @@ def reset_random():
         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),
         ]
 
         B0 = 4