Repository: pytorch/functorch (language: python)

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




Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesRandomness.cpp b/functorch/csrc/BatchRulesRandomness.cpp
index ac0d14fb8..3763a612d 100644
--- a/functorch/csrc/BatchRulesRandomness.cpp
+++ b/functorch/csrc/BatchRulesRandomness.cpp
@@ -35,6 +35,26 @@ Tensor random_batching_rule(IntArrayRef shape, ExtraArgs... extra_args) {
   }
 }
 
+template <typename F, F Func, typename... ExtraArgs>
+Tensor randperm_batching_rule(int64_t n, ExtraArgs... extra_args) {
+  c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  auto const batch_size = maybe_layer->batchSize();
+  RandomnessType randomness = maybe_layer->randomness();
+  check_randomness(randomness);
+  if (randomness == RandomnessType::Different) {
+    std::vector<at::Tensor> stackedList(batch_size);
+    stackedList.reserve(batch_size);
+    for (int64_t idx = 0; idx < batch_size; ++idx) {
+      // since this is done in a loop, need to pass by reference for generator to update
+      stackedList[idx] = Func(n, extra_args...);
+    }
+    return makeBatched(at::stack(stackedList), 0, maybe_layer->layerId());
+  } else {
+    return Func(n, std::forward<ExtraArgs>(extra_args)...);
+  }
+}
+
 template <typename A, A a, typename C>
 struct RandomBatchRuleHelper;
 
@@ -79,6 +99,16 @@ struct RandIntLowBatchRuleHelper<F, Func, typelist<T1, T2, T3, T...>> {
   }
 };
 
+template <typename A, A a, typename C>
+struct RandpermBatchRuleHelper;
+
+template <typename F, F Func, typename T1, typename... T>
+struct RandpermBatchRuleHelper<F, Func, typelist<T1, T...>> {
+  static Tensor apply(int64_t n, T... extra_args) {
+    return randperm_batching_rule<F, Func, T...>(n, std::forward<T>(extra_args)...);
+  }
+};
+
 TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   #define RANDOM_BATCH_RULE(op) \
     m.impl(#op, SINGLE_ARG(\
@@ -104,6 +134,15 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
     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))
+  #define RANDPERM_BATCH_RULE(op) \
+    m.impl(#op, SINGLE_ARG(\
+      RandpermBatchRuleHelper<decltype(&ATEN_FN(op)), &ATEN_FN(op), \
+                            c10::guts::function_traits<decltype(ATEN_FN(op))>::parameter_types>::apply))
+
+  #define RANDPERM_BATCH_RULE2(op, overload) \
+    m.impl(#op"."#overload, SINGLE_ARG(\
+      RandpermBatchRuleHelper<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);
@@ -120,10 +159,15 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   RANDINT_LOW_BATCH_RULE(randint, low);
   RANDINT_LOW_BATCH_RULE(randint, low_generator);
 
+  RANDPERM_BATCH_RULE(randperm);
+  RANDPERM_BATCH_RULE2(randperm, generator);
+
   #undef RANDOM_BATCH_RULE
   #undef RANDOM_BATCH_RULE2
   #undef RANDINT_BATCH_RULE
   #undef RANDINT_BATCH_RULE2
   #undef RANDINT_LOW_BATCH_RULE
+  #undef RANDPERM_BATCH_RULE
+  #undef RANDPERM_BATCH_RULE2
 }
 }} // namespace at::functorch
diff --git a/functorch/csrc/VmapModeRegistrations.cpp b/functorch/csrc/VmapModeRegistrations.cpp
index 893697675..ce3f4d01c 100644
--- a/functorch/csrc/VmapModeRegistrations.cpp
+++ b/functorch/csrc/VmapModeRegistrations.cpp
@@ -66,11 +66,6 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM(randint_like);
   UNSUPPORTED_RANDOM2(randint_like, low_dtype);
 
-  UNSUPPORTED_RANDOM(randperm);
-  UNSUPPORTED_RANDOM2(randperm, generator);
-  UNSUPPORTED_RANDOM2(randperm, out);
-  UNSUPPORTED_RANDOM2(randperm, generator_out);
-
   UNSUPPORTED_RANDOM(uniform_);
 }
 


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index b28abcd99..d094a71d8 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2647,9 +2647,6 @@ def test_no_random_op_support(self):
             (lambda t: captured.random_(0, 2), (torch.randn(B0),)),
             (lambda t: captured.random_(2), (torch.randn(B0),)),
             (lambda t: captured.uniform_(), (torch.randn(B0),)),
-
-            # factory functions
-            (lambda t: torch.randperm(5), (torch.randn(B0),)),
         ]
         for op, args in random_ops:
             with self.assertRaisesRegex(RuntimeError,
@@ -3462,6 +3459,37 @@ def reset_random():
                 for i in range(B0):
                     assert torch.allclose(vmap_result[i], expected)
 
+    @parametrize('randomness', ['same', 'different', 'error'])
+    @parametrize('use_generator', [True, False])
+    def test_randperm(self, device, randomness, use_generator):
+        # needs a special case because randperm doesn't take a batch size
+        B0 = 4
+        seed = 1234567
+        passed = torch.randn(B0, device=device)
+
+        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}
+
+        if randomness == 'error':
+            with self.assertRaisesRegex(RuntimeError, r"called random operation while in randomness error mode"):
+                vmap(lambda _: torch.randperm(10, **kwargs), randomness=randomness)(passed)
+            return
+
+        vmap_result = vmap(lambda _: torch.randperm(10, **kwargs), randomness=randomness)(passed)
+        generator = generator.set_state(orig_state)
+        torch.manual_seed(seed)
+        if randomness == 'different':
+            for i in range(B0):
+                expected = torch.randperm(10, **kwargs)
+                assert torch.allclose(vmap_result[i], expected)
+        else:
+            expected = torch.randperm(10, **kwargs)
+            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)