code: def check_dropout_replicated_input():
      dropout = nn.Dropout().cuda()
      dropout_replica = DropoutForReplicatedInput.from_native_module(
          dropout, process_group=None
      )
  
      # check computation correctness
      x = torch.rand(4, 128).cuda()
      out_1d = dropout_replica(x)
  
      # ensure out_1d is different across ranks
      world_size = dist.get_world_size()
      out_1d_all = [torch.zeros_like(out_1d) for _ in range(world_size)]
      dist.all_gather(out_1d_all, out_1d)
      for i in range(1, world_size):
          assert_equal(out_1d_all[i], out_1d_all[0])
  
file_path: colossalai/tests/test_shardformer/test_layer/test_dropout.py
imports: import torch
  import torch.distributed
  import torch.nn
  import colossalai
  from colossalai.shardformer.layer import DropoutForParallelInput, DropoutForReplicatedInput
  from colossalai.testing import assert_equal, assert_not_equal, rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 2 for calling a global function with example usage: check_dropout_replicated_input() without return types