Repository: pytorch/functorch (language: python)

Issue:
bernoulli_ only taking in float, includes dropout
Stack from [ghstack](https://github.com/ezyang/ghstack):
* __->__ #523
* #521




Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesDecompositions.cpp b/functorch/csrc/BatchRulesDecompositions.cpp
index c08ce23e2..539e8fe4f 100644
--- a/functorch/csrc/BatchRulesDecompositions.cpp
+++ b/functorch/csrc/BatchRulesDecompositions.cpp
@@ -84,6 +84,7 @@ void decompose_functional(const c10::OperatorHandle& op, torch::jit::Stack* stac
   }
 }
 
+
 #define DECOMPOSE_FUNCTIONAL(op) \
   m.impl(#op, torch::CppFunction::makeFromBoxedFunction<&decompose_functional>());
 
@@ -91,6 +92,17 @@ void decompose_functional(const c10::OperatorHandle& op, torch::jit::Stack* stac
 #define OP_DECOMPOSE(op)  m.impl(#op, static_cast<decltype(&ATEN_FN(op))>(native::op));
 #define OP_DECOMPOSE2(op, overload)  m.impl(#op"."#overload, static_cast<decltype(&ATEN_FN2(op, overload))>(native::op));
 
+TORCH_LIBRARY_IMPL(aten, FT_VMAP_MODE_KEY, m) {
+  OP_DECOMPOSE(alpha_dropout);
+  OP_DECOMPOSE(alpha_dropout_);
+  OP_DECOMPOSE(dropout);
+  OP_DECOMPOSE(dropout_);
+  OP_DECOMPOSE(feature_alpha_dropout);
+  OP_DECOMPOSE(feature_alpha_dropout_);
+  OP_DECOMPOSE(feature_dropout);
+  OP_DECOMPOSE(feature_dropout_);
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   OP_DECOMPOSE2(__and__, Scalar);
   OP_DECOMPOSE2(__and__, Tensor);
diff --git a/functorch/csrc/BatchRulesRandomness.cpp b/functorch/csrc/BatchRulesRandomness.cpp
index 2de8590be..629be1fe7 100644
--- a/functorch/csrc/BatchRulesRandomness.cpp
+++ b/functorch/csrc/BatchRulesRandomness.cpp
@@ -149,6 +149,39 @@ Tensor normal_batching_rule(const Tensor& tensor, double scalar, ExtraArgs... ex
   return makeBatched(out, 0, cur_level);
 }
 
+std::tuple<Tensor,Tensor> native_dropout_batching_rule(const Tensor& tensor, double p, c10::optional<bool> train) {
+  c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  const 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_value = moveBatchDimToFront(tensor_value, tensor_bdim);
+
+  if (!train.has_value() || train) {
+    check_randomness(randomness); // if we are in eval mode, we don't use about randomness
+  }
+
+  if ((train.has_value() && !train) || randomness == RandomnessType::Different) {
+    if (!tensor_bdim) {
+      tensor_value = tensor_value.unsqueeze(0).expand(maybe_layer->batchSize()); // ret value will be a batched tensor
+    }
+    auto res = at::native_dropout(tensor_value, p, train);
+    return std::make_tuple(makeBatched(std::get<0>(res), 0, cur_level), makeBatched(std::get<1>(res), 0, cur_level));
+  }
+
+  // repeated code from the CPU kernel since the CUDA one doesn't call bernoulli_ explicitly
+  double p1m = 1. - p;
+  // Check for probability of zero to avoid divide by zero and NaN results
+  double scale = p1m == 0 ? 0. : 1. / p1m;
+  Tensor mask = at::empty_like(tensor, LEGACY_CONTIGUOUS_MEMORY_FORMAT);
+  mask.bernoulli_(p1m);
+  const auto output = tensor.mul(mask).mul_(scale);
+  return std::make_tuple(output, mask);
+}
+
 template <typename A, A a, typename C>
 struct RandomBatchRuleHelper;
 
@@ -240,6 +273,17 @@ struct UnaryPointwiseRandomLeadingFloatBatchRule<F, Func, typelist<A0, A1, T...>
   }
 };
 
+TORCH_LIBRARY_IMPL(aten, FuncTorchBatched, m) {
+  #define RANDOM_INPLACE_BATCH_RULE2(op, overload) \
+    m.impl(#op"."#overload, SINGLE_ARG(\
+      RandomInplaceBatchRuleHelper<decltype(&ATEN_FN2(op, overload)), &ATEN_FN2(op, overload), \
+                            c10::guts::function_traits<decltype(ATEN_FN2(op, overload))>::parameter_types>::apply))
+  
+  RANDOM_INPLACE_BATCH_RULE2(bernoulli_, float);
+
+  #undef RANDOM_INPLACE_BATCH_RULE2
+}
+
 TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   #define RANDOM_BATCH_RULE(op) \
     m.impl(#op, SINGLE_ARG(\
@@ -317,6 +361,7 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   RAND_TWO_LEADING_SCALARS_BATCH_RULE(randint, low_generator);
 
   m.impl("bernoulli_.Tensor", at::functorch::bernoulli_inplace_Tensor_batching_rule);
+  RANDOM_INPLACE_BATCH_RULE2(bernoulli_, float);
 
   RANDPERM_BATCH_RULE(randperm);
   RANDPERM_BATCH_RULE2(randperm, generator);
@@ -325,6 +370,8 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNARY_POINTWISE_RANDOM(normal, Tensor_float);
   UNARY_POINTWISE_RANDOM_LEADING_FLOAT(normal, float_Tensor);
 
+  m.impl("native_dropout", native_dropout_batching_rule); // needs special casing because cuda version doesn't call bernoulli
+
   #undef RANDOM_BATCH_RULE
   #undef RANDOM_BATCH_RULE2
   #undef RANDOM_INPLACE_BATCH_RULE
diff --git a/functorch/csrc/VmapModeRegistrations.cpp b/functorch/csrc/VmapModeRegistrations.cpp
index 01ee05d45..4ec862f15 100644
--- a/functorch/csrc/VmapModeRegistrations.cpp
+++ b/functorch/csrc/VmapModeRegistrations.cpp
@@ -37,7 +37,6 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   UNSUPPORTED_RANDOM(bernoulli);
   UNSUPPORTED_RANDOM2(bernoulli, out);
   UNSUPPORTED_RANDOM2(bernoulli, p);
-  UNSUPPORTED_RANDOM(bernoulli_.float);
 
   UNSUPPORTED_RANDOM(cauchy_);
   UNSUPPORTED_RANDOM(exponential_);


Test patch:
diff --git a/test/test_eager_transforms.py b/test/test_eager_transforms.py
index 9990aa645..dffbe3332 100644
--- a/test/test_eager_transforms.py
+++ b/test/test_eager_transforms.py
@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import copy
 from torch.testing._internal.common_utils import (
     TestCase, run_tests, parametrize, subtest
 )
@@ -2319,6 +2320,71 @@ def slice_weights(batched_weights, index):
         self.assertEqual(result_loss, expected_loss)
         self.assertEqual(result_weights, expected_weights)
 
+    @parametrize("dropout_layer", [nn.Dropout, nn.AlphaDropout, nn.FeatureAlphaDropout])
+    def test_find_learning_rate_ensembling(self, device, dropout_layer):
+        # This example mimics what a user might do when trying to find the optimal learning rate. They would
+        # want to run a bunch of models with the same behavior (including the same dropout!) and have them
+        # each run with different learning rates. Specifically, this is an example of using same randomness with vmap
+        points, labels = torch.randn(100, 2, 2, 2, 2, device=device), torch.randint(0, 2, (100,), device=device)
+
+        class MLPClassifier(nn.Module):
+            def __init__(self, hidden_dim=32, n_classes=2):
+                super().__init__()
+                self.hidden_dim = hidden_dim
+                self.n_classes = n_classes
+
+                self.dropout = dropout_layer()
+                self.fc1 = nn.Linear(16, self.hidden_dim)
+                self.fc2 = nn.Linear(self.hidden_dim, self.n_classes)
+
+            def forward(self, x):
+                x = self.dropout(x)
+                x = torch.flatten(x, start_dim=1)
+                x = self.fc1(x)
+                x = F.relu(x)
+                x = self.fc2(x)
+                x = F.log_softmax(x, -1)
+                return x
+
+        loss_fn = nn.NLLLoss()
+
+        func_model, weights = make_functional(MLPClassifier().to(device))
+
+        def train_step_fn(weights, batch, targets, lr):
+            def compute_loss(weights, batch, targets):
+                output = func_model(weights, batch)
+                loss = loss_fn(output, targets)
+                return loss
+
+            grad_weights, loss = grad_and_value(compute_loss)(weights, batch, targets)
+            new_weights = []
+            with torch.no_grad():
+                for grad_weight, weight in zip(grad_weights, weights):
+                    new_weights.append(weight - grad_weight * lr)
+            # NB: return looks weird because torch.vmap must return Tensors
+            return (loss, *new_weights)
+
+        def unpack(train_result):
+            return train_result[0], train_result[1:]
+
+        def init_fn(num_models):
+            og_model = MLPClassifier().to(device)
+            models = tuple(copy.deepcopy(og_model) for _ in range(num_models))  # have same initialization
+            weights = tuple(make_functional(model)[1] for model in models)
+            weights = tuple(zip(*weights))
+            weights = tuple(torch.stack(shards).detach() for shards in weights)
+            return weights
+
+        batched_weights = init_fn(num_models=2)
+        parallel_train_step_fn = vmap(train_step_fn, in_dims=(0, None, None, 0), randomness="same")
+
+        lrs = torch.tensor([0.2, 0.4], device=device)
+        result_loss, result_weights = unpack(parallel_train_step_fn(batched_weights, points, labels, lrs))
+
+        self.assertEqual(result_loss[0], result_loss[1])
+        self.assertNotEqual(tuple(weight[0] for weight in result_weights),
+                            tuple(weight[1] for weight in result_weights))
+
     @unittest.skipIf(not USE_TORCHVISION, "test requires torchvision")
     def test_resnet18_per_sample_grads(self, device):
         import torchvision.models as models
diff --git a/test/test_vmap.py b/test/test_vmap.py
index bb6e7d43e..b877e5d94 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2656,7 +2656,6 @@ def test_no_random_op_support(self):
             (lambda t: torch.randint_like(captured, 0, 2), (torch.rand(B0),)),
 
             # in-place on BatchedTensor
-            (lambda t: t.bernoulli_(), (torch.randn(B0, 1),)),
             (lambda t: t.cauchy_(), (torch.randn(B0, 1),)),
             (lambda t: t.exponential_(), (torch.randn(B0, 1),)),
             (lambda t: t.geometric_(0.5), (torch.randn(B0, 1),)),
@@ -2664,7 +2663,6 @@ def test_no_random_op_support(self):
             (lambda t: t.uniform_(), (torch.randn(B0, 1),)),
 
             # in-place on captured tensor
-            (lambda t: captured.bernoulli_(), (torch.randn(B0),)),
             (lambda t: captured.cauchy_(), (torch.randn(B0),)),
             (lambda t: captured.exponential_(), (torch.randn(B0),)),
             (lambda t: captured.geometric_(0.5), (torch.randn(B0),)),
@@ -3114,7 +3112,7 @@ class TestVmapOperatorsOpInfo(TestCase):
         xfail('resize_as_'),
         xfail('tensor_split'),
         xfail('to_sparse'),
-        xfail('nn.functional.dropout'),
+        xfail('nn.functional.dropout'),  # works, can't check against for loop because of randomness inconsistency
         xfail('view_as_complex'),
         xfail('masked_select'),
 
@@ -3223,7 +3221,7 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('any'),
         xfail('count_nonzero'),
         xfail('nanmean'),
-        xfail('nn.functional.dropout'),
+        xfail('nn.functional.dropout'),  # works, can't check against for loop because of randomness inconsistency
         xfail('resize_'),
         xfail('view_as_complex'),
         xfail('matrix_exp'),
@@ -3489,6 +3487,7 @@ def test_random_behavior(self, device, randomness, use_generator):
             lambda _, shape: torch.normal(0., 1., shape, **kwargs),
             lambda t, _: t.normal_(**only_gen_kwarg),
             lambda t, _: t.bernoulli_(torch.tensor([0.3, 0.4, 0.5, 0.6]), **only_gen_kwarg),
+            lambda t, _: t.bernoulli_(**only_gen_kwarg),
         ]
 
         B0 = 4
@@ -3501,18 +3500,18 @@ def test_random_behavior(self, device, randomness, use_generator):
                     vmap(op, in_dims=(0, None), randomness=randomness)(passed, [B0])
                 return
 
-            passed = torch.randn(B0, B0, device=device)
+            passed = torch.randn(B0, B0, B0, device=device)
             generator = self.reset_random(generator, orig_state, use_generator, seed)
-            vmap_result = vmap(op, in_dims=(0, None), randomness=randomness)(passed, [B0])
+            vmap_result = vmap(op, in_dims=(0, None), randomness=randomness)(passed, [B0, B0])
             if randomness == "different":
-                passed = torch.randn([B0, B0], device=device)  # reset for in place operation
+                passed = torch.randn([B0, B0, B0], device=device)  # reset for in place operation
                 generator = self.reset_random(generator, orig_state, use_generator, seed)
-                expected = op(passed, [B0, B0])
+                expected = op(passed, [B0, B0, B0])
                 assert torch.allclose(vmap_result, expected)
             else:
-                passed = torch.randn(B0, device=device)  # reset for in place operation
+                passed = torch.randn(B0, B0, device=device)  # reset for in place operation
                 generator = self.reset_random(generator, orig_state, use_generator, seed)
-                expected = op(passed, [B0])
+                expected = op(passed, [B0, B0])
                 for i in range(B0):
                     assert torch.allclose(vmap_result[i], expected)
 
@@ -3547,6 +3546,38 @@ def test_randperm(self, device, randomness, use_generator):
             for i in range(B0):
                 assert torch.allclose(vmap_result[i], expected)
 
+    @parametrize('randomness', ['error', 'same', 'different'])
+    def test_feature_dropout(self, device, randomness):
+        # special case because functional feature dropout expects a batch dimension, so it doesn't match the
+        # pattern for other randomness
+
+        seed = 1234567
+        supported_ops = [
+            lambda t: torch.nn.functional.dropout(torch.ones_like(t), training=True),
+            lambda t: torch.nn.functional.alpha_dropout(torch.ones_like(t), training=True),
+            lambda t: torch.nn.functional.feature_alpha_dropout(t, training=True),
+            lambda t: torch.nn.functional.dropout2d(t, training=True),
+            lambda t: torch.nn.functional.dropout3d(t.unsqueeze(-1), training=True),
+        ]
+
+        for op in supported_ops:
+            passed = torch.ones([4, 3, 3, 3, 3])
+            if randomness == 'error':
+                with self.assertRaisesRegex(RuntimeError, r"called random operation while in randomness error mode"):
+                    vmap(op, randomness=randomness)(passed)
+                return
+            torch.manual_seed(seed)
+            vmap_result = vmap(op, randomness=randomness)(passed)
+            torch.manual_seed(seed)
+            if randomness == 'different':
+                expected = op(passed.flatten(0, 1))
+                expected = expected.reshape(passed.shape[0], passed.shape[1], *expected.shape[1:])
+                assert torch.allclose(vmap_result, expected)
+            else:
+                expected = op(passed[0])
+                for i in range(4):
+                    assert torch.allclose(vmap_result[i], expected)
+
     @parametrize('use_generator', [True, False])
     @parametrize('randomness', ['error', 'same', 'different'])
     def test_random_inplace_not_batched(self, device, use_generator, randomness):
@@ -3557,7 +3588,8 @@ def test_random_inplace_not_batched(self, device, use_generator, randomness):
             lambda t, _: t.random_(100, **kwargs),
             lambda t, _: t.random_(-5, 100, **kwargs),
             lambda t, _: t.normal_(**kwargs),
-            lambda t, _: t.bernoulli_(torch.tensor([0.3, 0.4, 0.5, 0.6]), **kwargs)
+            lambda t, _: t.bernoulli_(torch.tensor([0.3, 0.4, 0.5, 0.6]), **kwargs),
+            lambda t, _: t.bernoulli_(**kwargs),
         ]
 
         B0 = 4