add unit test for other_ops.py

KingsleyLiu-NV · KingsleyLiu-NV · commit 2e52775e7470 · 2023-11-19T21:33:28.000-08:00
diff --git a/test/distributed/_tensor/test_other_ops.py b/test/distributed/_tensor/test_other_ops.py
@@ -6,14 +6,14 @@
 
 import torch.distributed as dist
 
-from timm.models.layers import DropPath
 from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate
 from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     DTensorTestBase,
     with_comms,
 )
 
+
 ITER_TIME = 10
 LR = 0.001
 
@@ -84,18 +84,16 @@ def test_bernoulli(self):
         input_list = torch.rand(ITER_TIME, 1024, 10)
         grad_output_list = torch.rand(ITER_TIME, 1024, 10) * 1e-3
 
-        droppath = DropPath(0.3)
-
         for i in range(ITER_TIME):
             inp = input_list[i].to(self.device_type).requires_grad_()
             grad_output = grad_output_list[i].to(self.device_type)
 
-            # droppath  with dtensor
+            # bernoulli  with dtensor
             inp_dtensor = distribute_tensor(inp, device_mesh, shard_spec)
             grad_output_dtensor = distribute_tensor(
                 grad_output, device_mesh, shard_spec
             )
-            output = droppath(inp_dtensor)
+            output = torch.bernoulli(inp_dtensor)
             output.backward(grad_output_dtensor)
 
             send_output_tensor = output.to_local()
diff --git a/torch/distributed/_tensor/ops/other_ops.py b/torch/distributed/_tensor/ops/other_ops.py
@@ -36,6 +36,7 @@ def slice_backward_rules(op_schema: OpSchema) -> OutputSharding:
     return OutputSharding(grad_input_spec)
 
 
+@register_prop_rule(aten.bernoulli.default)
 @register_prop_rule(aten.bernoulli_.float)
 def bernoulli_rules(op_schema: OpSchema) -> OutputSharding:
     input_spec = op_schema.args_schema[0]
