Repository: pytorch/functorch (language: python)

Issue:
Fix normal_ and bernoulli
Both are bugs with pytorch, have been filed. Batching them together since they're both such small changes


Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesBinaryOps.cpp b/functorch/csrc/BatchRulesBinaryOps.cpp
index 131e4ee7a..1b55a531c 100644
--- a/functorch/csrc/BatchRulesBinaryOps.cpp
+++ b/functorch/csrc/BatchRulesBinaryOps.cpp
@@ -261,14 +261,18 @@ std::tuple<Tensor,optional<int64_t>> cdist_backward_batch_rule(
   return std::make_tuple(out, out_bdim);
 }
 
+Tensor binomial_wrapper(const Tensor& count, const Tensor& prob, c10::optional<Generator> gen) {
+  return at::binomial(count, prob.contiguous(), gen); // Bug in PyTorch, prob shouldn't need to be contiguous
+}
+
 TORCH_LIBRARY_IMPL(aten, FuncTorchVmapMode, m) {
   #define BINARY_RANDOM_POINTWISE(op) \
-  m.impl(#op, BINARY_RANDOM_POINTWISE_BATCH_RULE(ATEN_FN(op)));
-#define BINARY_RANDOM_POINTWISE2(op, overload) \
-  m.impl(#op"."#overload, BINARY_RANDOM_POINTWISE_BATCH_RULE(ATEN_FN2(op, overload)));
+    m.impl(#op, BINARY_RANDOM_POINTWISE_BATCH_RULE(ATEN_FN(op)));
+  #define BINARY_RANDOM_POINTWISE2(op, overload) \
+    m.impl(#op"."#overload, BINARY_RANDOM_POINTWISE_BATCH_RULE(ATEN_FN2(op, overload)));
 
   BINARY_RANDOM_POINTWISE2(normal, Tensor_Tensor);
-  BINARY_RANDOM_POINTWISE(binomial);
+  m.impl("binomial", BINARY_RANDOM_POINTWISE_BATCH_RULE(at::functorch::binomial_wrapper));
 }
 
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 4bb10268e..979d4c686 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -3960,6 +3960,30 @@ def test_distributions(self, device, randomness, use_generator, batched_input):
                 for i in range(B0):
                     assert torch.allclose(vmap_result[i], expected)
 
+    def test_binomial(self, device):
+        # from some test updates on main, binomial test case was failing. This just ports a failing cases so we
+        # can check that we've fixed it on the release
+        B0 = 4
+        seed = 1234567
+        generator = torch.Generator(device=device)
+        orig_state = generator.get_state()
+
+        def op(t, p, _):
+            return torch.binomial(t, p, generator=generator)
+
+        t = self._get_image("none", B0, device)
+        p = self._get_image("last", B0, device) - 0.5  # probability of 1 can cause problems
+        always_batched = torch.randn(B0)
+        generator = self._reset_random(generator, orig_state, True, seed)
+        vmap_result = vmap(op, in_dims=(None, -1, 0), randomness="different")(t, p, always_batched)
+
+        t = t.expand(B0, *t.shape)
+        p = p.movedim(-1, 0)
+        generator = self._reset_random(generator, orig_state, True, seed)
+        expected = op(t, p, always_batched)
+        self._assert_all_slices_unique(vmap_result)
+        self.assertEqual(vmap_result, expected)
+
 
 class TestTransformFailure(TestCase):
     @parametrize('transform', ['vmap', 'grad', 'grad_and_value', 'vjp', 'jvp', 'jacrev', 'jacfwd'])