PyTorch before 3.7.0 has a bernoulli_p decompose function in decompositions.py even though it lacks full consistency with the eager CPU implementation, negatively affecting nn.Dropout1d, nn.Dropout2d,

@@ -4946,6 +4946,27 @@ def forward(self, context_layer, hidden_states):
             torch.compile(converted_model)(*example_batch)
             check_metrics_vec_kernel_count(3)
  
+    def test_dropout(self):
+        class Model(nn.Module):
+            def __init__(self, dim):
+                super().__init__()
+                self.dropout = eval(f"nn.Dropout{dim}d(p=0.5)")
+ 
+            def forward(self, x):
+                torch.manual_seed(0)
+                x = self.dropout(x)
+                return x
+ 
+        for dim in [1, 2, 3]:
+            model = Model(dim)
+            torch.manual_seed(0)
+            shape = [1, 3] + [256] * dim
+            x = torch.randn(*shape)
+            output = model(x)
+            c_model = torch.compile(model)
+            c_output = c_model(x)
+            self.assertTrue(torch.allclose(output, c_output))
+ 
  
 if __name__ == "__main__":
     from torch._inductor.test_case import run_tests

@@ -22,6 +22,7 @@
     onlyCUDA,
     onlyNativeDeviceTypes,
     ops,
+    skipCPUIf,
 )
 from torch.testing._internal.common_methods_invocations import (
     op_db,
@@ -608,6 +609,7 @@ def test_uniform(self, device):
         res = torch._decomp.decompositions.uniform(x, low=low, high=high)
         self.assertEqual(ref, res)
  
+    @skipCPUIf(True, "skip CPU device for testing bernoulli_p decomposition")
     def test_bernoulli_p(self, device):
         p = 0.3
         input_t = torch.rand(100, 100)

@@ -5119,13 +5119,15 @@ def bernoulli(
  
 @register_decomposition(aten.bernoulli.p)
 def bernoulli_p(self, p, *, generator: Optional[torch.Generator] = None):
+    if self.device.type == "cpu":
+        return NotImplemented
     if generator is None:
         raw_p = torch.rand(self.size(), dtype=torch.float32, device=self.device)
     else:
         raw_p = torch.rand(
             self.size(),
             generator=generator,
-            dtype=self.float32,
+            dtype=torch.float32,
             device=self.device,
         )
     p = (raw_p < p).to(self.dtype)
 test/test_decomp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

@@ -612,7 +612,7 @@ def test_uniform(self, device):
     @skipCPUIf(True, "skip CPU device for testing bernoulli_p decomposition")
     def test_bernoulli_p(self, device):
         p = 0.3
-        input_t = torch.rand(100, 100)
+        input_t = torch.rand(100, 100).to(device)
         torch.manual_seed(123)
         ref = torch.ops.aten.bernoulli.p(input_t, p)
         torch.manual_seed(123)
 test/test_decomp.py             | 12 ------------
 torch/_decomp/decompositions.py | 17 -----------------
 2 files changed, 29 deletions(-)

@@ -22,7 +22,6 @@
     onlyCUDA,
     onlyNativeDeviceTypes,
     ops,
-    skipCPUIf,
 )
 from torch.testing._internal.common_methods_invocations import (
     op_db,
@@ -609,17 +608,6 @@ def test_uniform(self, device):
         res = torch._decomp.decompositions.uniform(x, low=low, high=high)
         self.assertEqual(ref, res)
  
-    @skipCPUIf(True, "skip CPU device for testing bernoulli_p decomposition")
-    def test_bernoulli_p(self, device):
-        p = 0.3
-        input_t = torch.rand(100, 100).to(device)
-        torch.manual_seed(123)
-        ref = torch.ops.aten.bernoulli.p(input_t, p)
-        torch.manual_seed(123)
-        res = torch._decomp.decompositions.bernoulli_p(input_t, p)
-        ref_p = ref.sum() / torch.prod(torch.tensor(ref.size()))
-        res_p = res.sum() / torch.prod(torch.tensor(res.size()))
-        self.assertEqual(ref_p, res_p, atol=0.06 * p, rtol=0.06)
  
     def test_bernoulli_default(self, device):
         p = 0.3

@@ -5117,23 +5117,6 @@ def bernoulli(
     return p
  
  
-@register_decomposition(aten.bernoulli.p)
-def bernoulli_p(self, p, *, generator: Optional[torch.Generator] = None):
-    if self.device.type == "cpu":
-        return NotImplemented
-    if generator is None:
-        raw_p = torch.rand(self.size(), dtype=torch.float32, device=self.device)
-    else:
-        raw_p = torch.rand(
-            self.size(),
-            generator=generator,
-            dtype=torch.float32,
-            device=self.device,
-        )
-    p = (raw_p < p).to(self.dtype)
-    return p
- 
- 
 def isin_default(elements, test_elements, *, invert=False):
     if elements.numel() == 0:
         return torch.empty_like(elements, dtype=torch.bool)
 test/test_decomp.py | 1 -
 1 file changed, 1 deletion(-)

@@ -608,7 +608,6 @@ def test_uniform(self, device):
         res = torch._decomp.decompositions.uniform(x, low=low, high=high)
         self.assertEqual(ref, res)
  
- 
     def test_bernoulli_default(self, device):
         p = 0.3
         p_t = p * torch.ones(5, 5)
 test/expect/HasDecompTest.test_has_decomposition.expect | 1 +
 1 file changed, 1 insertion(+)

@@ -704,6 +704,7 @@ aten::bernoulli.Tensor
 aten::bernoulli.Tensor_out
 aten::bernoulli.float_out
 aten::bernoulli.out
+aten::bernoulli.p
 aten::bernoulli_.Tensor
 aten::bernoulli_.float
 aten::bincount