Repository: pytorch/functorch (language: python)

Issue:
bernoulli_ support
Stack from [ghstack](https://github.com/ezyang/ghstack):
* #523
* __->__ #521




Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesRandomness.cpp b/functorch/csrc/BatchRulesRandomness.cpp
index 6da49b276..2de8590be 100644
--- a/functorch/csrc/BatchRulesRandomness.cpp
+++ b/functorch/csrc/BatchRulesRandomness.cpp
@@ -53,6 +53,55 @@ Tensor& random_inplace_batching_rule(Tensor& self, ExtraArgs... extra_args) {
   }
 }
 
+Tensor& bernoulli_inplace_Tensor_batching_rule(Tensor& self, const Tensor& p_, c10::optional<Generator> gen) {
+  c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
+  auto maybe_layer = maybeCurrentDynamicLayer();
+  auto cur_level = maybe_layer->layerId();
+  RandomnessType randomness = maybe_layer->randomness();
+
+  Tensor self_value;
+  optional<int64_t> self_bdim;
+  std::tie(self_value, self_bdim) = unwrapTensorAtLevel(self, cur_level);
+
+  Tensor other_value;
+  optional<int64_t> other_bdim;
+  std::tie(other_value, other_bdim) = unwrapTensorAtLevel(p_, cur_level);
+
+  check_randomness(randomness, other_bdim.has_value());
+
+  if (!self_bdim && other_bdim) {
+    vmapIncompatibleInplaceError("inplace arithmetic");
+  }
+
+  // compute max logical rank
+  auto self_logical_rank = rankWithoutBatchDim(self_value, self_bdim);
+  auto other_logical_rank = rankWithoutBatchDim(other_value, other_bdim);
+  auto max_logical_rank = std::max(self_logical_rank, other_logical_rank);
+
+  auto self_ = moveBatchDimToFront(self_value, self_bdim);
+  auto other_ = moveBatchDimToFront(other_value, other_bdim);
+
+  // If the dimensions aren't aligned, we need to line them up.
+  // Tensor[B, 3] + Tensor[2, 5, 3] -> Tensor[B, 1, 1, 3] + Tensor[2, 5, 3]
+  // Note that only tensors that have a batch dim need to be modified.
+  // Tensor[B, 2, 3, 5] + Tensor[5] -> no changes needed
+  self_ = maybePadToLogicalRank(self_, self_bdim, max_logical_rank);
+  other_ = maybePadToLogicalRank(other_, other_bdim, max_logical_rank);
+  TORCH_CHECK(
+    !(randomness == RandomnessType::Different && !self_bdim),
+    "vmap: Cannot ask for different inplace randomness on an unbatched tensor. This will appear like same randomness. ",
+    "If this is necessary for your usage, please file an issue with functorch.");
+  if (randomness == RandomnessType::Same && self_bdim) {
+    auto intermediate = empty(self.sizes(), self.options());
+    intermediate.bernoulli_(other_, gen);
+    self.copy_(intermediate); // batching should make this just work out...
+    return self;
+  } else {
+    self_.bernoulli_(other_, gen);
+    return self;
+  }
+}
+
 template <typename F, F Func, typename... ExtraArgs>
 Tensor randperm_batching_rule(int64_t n, ExtraArgs... extra_args) {
   c10::impl::ExcludeDispatchKeyGuard guard(kVmapModeKey);
@@ -267,6 +316,8 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   RAND_TWO_LEADING_SCALARS_BATCH_RULE(randint, low);
   RAND_TWO_LEADING_SCALARS_BATCH_RULE(randint, low_generator);
 
+  m.impl("bernoulli_.Tensor", at::functorch::bernoulli_inplace_Tensor_batching_rule);
+
   RANDPERM_BATCH_RULE(randperm);
   RANDPERM_BATCH_RULE2(randperm, generator);
 
diff --git a/functorch/csrc/VmapModeRegistrations.cpp b/functorch/csrc/VmapModeRegistrations.cpp
index da3790b92..01ee05d45 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_RANDOM2(bernoulli_, Tensor);
   UNSUPPORTED_RANDOM(bernoulli_.float);
 
   UNSUPPORTED_RANDOM(cauchy_);


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 4f63477dd..bb6e7d43e 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -3488,6 +3488,7 @@ def test_random_behavior(self, device, randomness, use_generator):
             lambda t, _: t.random_(**only_gen_kwarg),
             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),
         ]
 
         B0 = 4
@@ -3556,6 +3557,7 @@ 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)
         ]
 
         B0 = 4