Repository: pytorch/functorch (language: python)

Issue:
Add support for randn and rand
Stack from [ghstack](https://github.com/ezyang/ghstack):
* #420
* #412
* #411
* #410
* __->__ #409

Updates `vmap` to accept a randomness parameter ('error', 'different', or 'same') that determines what the randomness behavior should be. Also adds `torch.randn` and `torch.rand` to demonstrate basic examples


Discussion / hints:

Fix patch:
diff --git a/functorch/_src/vmap.py b/functorch/_src/vmap.py
index c2cdd86fd..c2f1ecd8b 100644
--- a/functorch/_src/vmap.py
+++ b/functorch/_src/vmap.py
@@ -203,7 +203,11 @@ def _get_name(func: Callable):
 # on BatchedTensors perform the batched operations that the user is asking for.
 
 
-def vmap(func: Callable, in_dims: in_dims_t = 0, out_dims: out_dims_t = 0) -> Callable:
+def vmap(
+        func: Callable,
+        in_dims: in_dims_t = 0,
+        out_dims: out_dims_t = 0,
+        randomness: str = 'error') -> Callable:
     """
     vmap is the vectorizing map; ``vmap(func)`` returns a new function that
     maps :attr:`func` over some dimension of the inputs. Semantically, vmap
@@ -226,6 +230,11 @@ def vmap(func: Callable, in_dims: in_dims_t = 0, out_dims: out_dims_t = 0) -> Ca
         out_dims (int or Tuple[int]): Specifies where the mapped dimension
             should appear in the outputs. If :attr:`out_dims` is a Tuple, then
             it should have one element per output. Default: 0.
+        use_batched_random (str): Specifies whether the randomness in this
+            vmap should be the same or different across batches. If 'different',
+            the randomness for each batch will be different. If 'same', the
+            randomness will be the same across batches. If 'error', any calls to
+            random functions will error. Default: True.
 
     Returns:
         Returns a new "batched" function. It takes the same inputs as
@@ -343,11 +352,14 @@ def vmap(func: Callable, in_dims: in_dims_t = 0, out_dims: out_dims_t = 0) -> Ca
         vmap does not provide general autobatching or handle variable-length
         sequences out of the box.
     """
+    if randomness not in ['error', 'different', 'same']:
+        raise RuntimeError(f"Only allowed values for randomness are 'error', 'different', or 'same'. Got {randomness}")
+
     @functools.wraps(func)
     def wrapped(*args, **kwargs):
         _check_out_dims_is_int_or_int_pytree(out_dims, func)
         batch_size, flat_in_dims, flat_args, args_spec = _process_batched_inputs(in_dims, args, func)
-        vmap_level = _vmap_increment_nesting(batch_size)
+        vmap_level = _vmap_increment_nesting(batch_size, randomness)
         try:
             batched_inputs = _create_batched_inputs(flat_in_dims, flat_args, vmap_level, args_spec)
             batched_outputs = func(*batched_inputs, **kwargs)
diff --git a/functorch/csrc/BatchRulesRandomness.cpp b/functorch/csrc/BatchRulesRandomness.cpp
new file mode 100644
index 000000000..02e27c34e
--- /dev/null
+++ b/functorch/csrc/BatchRulesRandomness.cpp
@@ -0,0 +1,73 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <ATen/ATen.h>
+#include <functorch/csrc/DynamicLayer.h>
+#include <functorch/csrc/BatchRulesHelper.h>
+
+namespace at {
+namespace functorch {
+
+void check_randomness(std::string randomness) {
+  TORCH_CHECK(
+    randomness != "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"
+  );
+}
+
+template <typename F, F Func, typename... ExtraArgs>
+Tensor random_batching_rule(IntArrayRef shape, ExtraArgs... extra_args) {
+  c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  VmapDimVector shapeVec(shape.begin(), shape.end());
+  shapeVec.insert(shapeVec.begin(), maybe_layer->batchSize());
+  std::string randomness = maybe_layer->randomness();
+  check_randomness(randomness);
+  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());
+  }
+}
+
+template <typename A, A a, typename C>
+struct RandomBatchRuleHelper;
+
+template <typename F, F Func, typename T1, typename... T>
+struct RandomBatchRuleHelper<F, Func, typelist<T1, T...>> {
+  static Tensor apply(IntArrayRef shape, T... extra_args) {
+    return random_batching_rule<F, Func, T...>(shape, std::forward<T>(extra_args)...);
+  }
+};
+
+
+TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
+  #define RANDOM_BATCH_RULE(op) \
+    m.impl(#op, SINGLE_ARG(\
+      RandomBatchRuleHelper<decltype(&ATEN_FN(op)), &ATEN_FN(op), \
+                            c10::guts::function_traits<decltype(ATEN_FN(op))>::parameter_types>::apply))
+
+  #define RANDOM_BATCH_RULE2(op, overload) \
+    m.impl(#op"."#overload, SINGLE_ARG(\
+      RandomBatchRuleHelper<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);
+  RANDOM_BATCH_RULE2(randn, names);
+
+  RANDOM_BATCH_RULE(rand);
+  RANDOM_BATCH_RULE2(rand, generator);
+  RANDOM_BATCH_RULE2(rand, generator_with_names);
+  RANDOM_BATCH_RULE2(rand, names);
+
+  #undef RANDOM_BATCH_RULE
+  #undef RANDOM_BATCH_RULE2
+}
+}} // namespace at::functorch
\ No newline at end of file
diff --git a/functorch/csrc/DynamicLayer.cpp b/functorch/csrc/DynamicLayer.cpp
index 3eae8ad79..3ba8e78b2 100644
--- a/functorch/csrc/DynamicLayer.cpp
+++ b/functorch/csrc/DynamicLayer.cpp
@@ -50,11 +50,13 @@ DynamicLayer::DynamicLayer(
     DispatchKey key,
     int64_t layerId,
     optional<int64_t> batchSize,
+    optional<std::string> randomness,
     optional<bool> prev_grad_mode)
   :
     key_(key),
     layerId_(layerId),
     batchSize_(batchSize),
+    randomness_(randomness),
     prevGradMode_(prev_grad_mode)
 {
   if (key_ == DispatchKey::Autograd) {
@@ -75,6 +77,11 @@ int64_t DynamicLayer::batchSize() const {
   return *batchSize_;
 }
 
+std::string DynamicLayer::randomness() const {
+  TORCH_INTERNAL_ASSERT(randomness_);
+  return *randomness_;
+}
+
 optional<bool> DynamicLayer::prevGradMode() const {
   return prevGradMode_;
 }
@@ -110,7 +117,7 @@ class FuncTorchTLS : public FuncTorchTLSBase {
 
   // Initial autograd layer, because autograd is always "on"
   // TODO: Get rid of this, it is bad for composability
-  std::vector<DynamicLayer> dynamicLayerStack = { DynamicLayer(DispatchKey::Autograd, 1, nullopt, true) };
+  std::vector<DynamicLayer> dynamicLayerStack = { DynamicLayer(DispatchKey::Autograd, 1, nullopt, nullopt, true) };
 };
 
 static FuncTorchTLS* getRawFunctorchTLS() {
@@ -191,11 +198,12 @@ static int64_t pushDynamicLayer(DynamicLayer&& dynamic_layer) {
 int64_t initAndPushDynamicLayer(
     DispatchKey key,
     optional<int64_t> batch_size,
+    optional<std::string> randomness,
     optional<bool> prev_grad_mode) {
   TORCH_INTERNAL_ASSERT(key == DispatchKey::Autograd || key == kBatchedKey);
   const auto& dynamicLayerStack = dynamicLayerStackAccessor();
   const auto layerId = 1 + dynamicLayerStack.size();
-  DynamicLayer new_layer(key, layerId, batch_size, prev_grad_mode);
+  DynamicLayer new_layer(key, layerId, batch_size, randomness, prev_grad_mode);
   pushDynamicLayer(std::move(new_layer));
 
   auto& data = getGlobalDynmetaData();
diff --git a/functorch/csrc/DynamicLayer.h b/functorch/csrc/DynamicLayer.h
index 88ff6e7ed..77cc6881f 100644
--- a/functorch/csrc/DynamicLayer.h
+++ b/functorch/csrc/DynamicLayer.h
@@ -22,6 +22,7 @@ struct TORCH_API DynamicLayer {
       DispatchKey key,
       int64_t layerId,
       optional<int64_t> batchSize = nullopt,
+      optional<std::string> randomness = nullopt,
       optional<bool> prev_grad_mode = nullopt);
 
   DispatchKey key() const;
@@ -29,6 +30,7 @@ struct TORCH_API DynamicLayer {
 
   // Only valid for vmap
   int64_t batchSize() const;
+  std::string randomness() const;
 
   // only valid for grad-based transforms
   optional<bool> prevGradMode() const;
@@ -39,12 +41,14 @@ struct TORCH_API DynamicLayer {
   // Honestly these should be a union or some extendable metadata class.
   // Not doing that for now because I don't think we'll use this mechanism for very long.
   optional<int64_t> batchSize_;
+  optional<std::string> randomness_;
   optional<bool> prevGradMode_;
 };
 
 TORCH_API int64_t initAndPushDynamicLayer(
     DispatchKey key,
     optional<int64_t> batch_size = nullopt,
+    optional<std::string> randomness = nullopt,
     optional<bool> prev_grad_mode = nullopt);
 TORCH_API DynamicLayer popDynamicLayerAndDeleteMetadata();
 TORCH_API c10::optional<DynamicLayer> maybeCurrentDynamicLayer();
diff --git a/functorch/csrc/VmapModeRegistrations.cpp b/functorch/csrc/VmapModeRegistrations.cpp
index bb455cc46..d2c0a8cc4 100644
--- a/functorch/csrc/VmapModeRegistrations.cpp
+++ b/functorch/csrc/VmapModeRegistrations.cpp
@@ -32,18 +32,6 @@ TORCH_LIBRARY_IMPL(_, FuncTorchVmapMode, m) {
 #define UNSUPPORTED_RANDOM2(op, overload) \
   m.impl(#op"."#overload, torch::CppFunction::makeFromBoxedFunction<&unsupportedRandomOp>());
 
-#define TENSOROPTIONSPARAMS c10::optional<c10::ScalarType> dtype, c10::optional<c10::Layout> layout, c10::optional<c10::Device> device, c10::optional<bool> pin_memory
-#define TENSOROPTIONSARGS dtype, layout, device, pin_memory
-
-Tensor randn_mbatching_rule(IntArrayRef shape, TENSOROPTIONSPARAMS) {
-    TORCH_WARN("Automatically expanding the shape of randn. This means that different elements of vmap will get different random values. If you want different behavior, like using the same random values across vmap, please make a github issue.");
-    c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
-    auto maybe_layer = maybeCurrentDynamicLayer();
-    VmapDimVector shapeVec(shape.begin(), shape.end());
-    shapeVec.insert(shapeVec.begin(), maybe_layer->batchSize());
-    return makeBatched(at::randn(shapeVec, TENSOROPTIONSARGS), 0, maybe_layer->layerId());
-}
-
 
 TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM(bernoulli);
@@ -75,26 +63,9 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM2(random_, to);
   UNSUPPORTED_RANDOM(random_);
 
-  // UNSUPPORTED_RANDOM(rand_like);
-  // UNSUPPORTED_RANDOM(randn_like);
-
   UNSUPPORTED_RANDOM(randint_like);
   UNSUPPORTED_RANDOM2(randint_like, low_dtype);
 
-  UNSUPPORTED_RANDOM(rand);
-  UNSUPPORTED_RANDOM2(rand, generator);
-  UNSUPPORTED_RANDOM2(rand, names);
-  UNSUPPORTED_RANDOM2(rand, generator_with_names);
-  UNSUPPORTED_RANDOM2(rand, out);
-  UNSUPPORTED_RANDOM2(rand, generator_out);
-
-//   UNSUPPORTED_RANDOM(randn);
-  UNSUPPORTED_RANDOM2(randn, generator);
-  UNSUPPORTED_RANDOM2(randn, names);
-  UNSUPPORTED_RANDOM2(randn, generator_with_names);
-  UNSUPPORTED_RANDOM2(randn, out);
-  UNSUPPORTED_RANDOM2(randn, generator_out);
-
   UNSUPPORTED_RANDOM(randperm);
   UNSUPPORTED_RANDOM2(randperm, generator);
   UNSUPPORTED_RANDOM2(randperm, out);
@@ -110,9 +81,6 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM2(randint, low_generator_out);
 
   UNSUPPORTED_RANDOM(uniform_);
-
-
-  m.impl("randn", randn_mbatching_rule);
 }
 
 
diff --git a/functorch/csrc/init.cpp b/functorch/csrc/init.cpp
index 3066113d2..dc2f5f70e 100644
--- a/functorch/csrc/init.cpp
+++ b/functorch/csrc/init.cpp
@@ -138,7 +138,7 @@ bool dump_tensor(const Tensor& self) {
 int64_t _grad_increment_nesting() {
   // See NOTE [grad and vjp interaction with no_grad]
   bool prev_grad_mode = c10::GradMode::is_enabled();
-  return initAndPushDynamicLayer(at::DispatchKey::Autograd, nullopt, prev_grad_mode);
+  return initAndPushDynamicLayer(at::DispatchKey::Autograd, nullopt, nullopt, prev_grad_mode);
 }
 
 int64_t _grad_decrement_nesting() {
@@ -147,8 +147,8 @@ int64_t _grad_decrement_nesting() {
   return layer.layerId();
 }
 
-int64_t _vmap_increment_nesting(int64_t batch_size) {
-  return initAndPushDynamicLayer(kBatchedKey, batch_size);
+int64_t _vmap_increment_nesting(int64_t batch_size, std::string randomness) {
+  return initAndPushDynamicLayer(kBatchedKey, batch_size, randomness);
 }
 
 int64_t _vmap_decrement_nesting() {


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
index ba4c2b3fa..220129aa1 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -905,10 +905,10 @@ def test():
         xfail('nn.functional.gaussian_nll_loss'),
         xfail('double', 'channels_last'),
         xfail('masked_select'),
-        xfail('nn.functional.fractional_max_pool3d'),
+        skip('nn.functional.fractional_max_pool3d'),  # generator works on cpu, fails on cuda
         xfail('nn.functional.glu'),
         xfail('as_strided'),
-        xfail('nn.functional.fractional_max_pool2d'),
+        skip('nn.functional.fractional_max_pool2d'),  # generator works on cpu, fails on cuda
         skip('solve'),
         xfail('linalg.cond'),
         xfail('linalg.svdvals'),
diff --git a/test/test_vmap.py b/test/test_vmap.py
index da996d55a..279efb8e5 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2592,7 +2592,6 @@ def test_view_as(self):
              (torch.rand(B1, B2, B0, 3, 2, 5), torch.rand(B0, 3 * 2 * 5)),
              in_dims=(2, 0))
 
-    # TODO: reenable the random op failures
     def test_no_random_op_support(self):
         B0 = 2
 
@@ -2650,8 +2649,6 @@ def test_no_random_op_support(self):
             (lambda t: captured.uniform_(), (torch.randn(B0),)),
 
             # factory functions
-            (lambda t: torch.rand(1), (torch.randn(B0),)),
-            # (lambda t: torch.randn(1), (torch.randn(B0),)),
             (lambda t: torch.randint(5, [1]), (torch.randn(B0),)),
             (lambda t: torch.randperm(5), (torch.randn(B0),)),
         ]
@@ -2716,14 +2713,14 @@ def naive_f(x, shape):
             return x + torch.randn(shape)
 
         torch.manual_seed(0)
-        out1 = vmap(vmap(vmap_f))(torch.ones(2, 3))
+        out1 = vmap(vmap(vmap_f, randomness='different'), randomness='different')(torch.ones(2, 3))
 
         torch.manual_seed(0)
         out2 = naive_f(torch.ones(2, 3), (2, 3))
         self.assertEqual(out1, out2)
 
         torch.manual_seed(0)
-        out1 = vmap(vmap(vmap_f))(torch.ones(2, 3, 4))
+        out1 = vmap(vmap(vmap_f, randomness='different'), randomness='different')(torch.ones(2, 3, 4))
 
         torch.manual_seed(0)
         out2 = naive_f(torch.ones(2, 3, 4), (2, 3, 1))
@@ -3429,6 +3426,42 @@ def f(x, y):
         y = torch.randn(2, 3, device=device)
         self.assertTrue(isinstance(vmap(f)(x, y), Point))
 
+    @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}
+        supported_random_ops = [
+            lambda _, shape: torch.randn(shape, **kwargs),
+            lambda _, shape: torch.rand(shape, **kwargs),
+        ]
+
+        B0 = 4
+        seed = 1234567
+        passed = torch.randn(B0, device=device)
+
+        for op in supported_random_ops:
+            if randomness == 'error':
+                with self.assertRaisesRegex(RuntimeError, r"called random operation while in randomness error mode"):
+                    vmap(op, in_dims=(0, None), randomness=randomness)(passed, [B0])
+                return
+
+            generator = reset_random()
+            vmap_result = vmap(op, in_dims=(0, None), randomness=randomness)(passed, [B0])
+
+            generator = reset_random()
+            if randomness == 'different':
+                expected = op(passed, [B0, B0])
+                assert torch.allclose(vmap_result, expected)
+            else:
+                expected = op(passed, [B0])
+                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)