def METHOD_NAME(device_mesh, rank):
    # tensor to comm
    sharded_tensor_to_comm_0 = torch.zeros(2, 2).cuda()
    sharded_tensor_to_comm_1 = torch.ones(2, 2).cuda()
    # tensor([[0., 0., 1., 1.],
    #         [0., 0., 1., 1.]])
    tensor_to_shard = torch.cat((sharded_tensor_to_comm_0, sharded_tensor_to_comm_1), 1)
    # test shard
    dim_partition_dict = {}
    # DistSpec:
    #     shard_sequence: R,R
    #     device_mesh_shape: (2, 2)
    sharding_spec = ShardingSpec(device_mesh, tensor_to_shard.shape, dim_partition_dict=dim_partition_dict)
    # CommSpec:(comm_pattern:shard, shard_dim:1, logical_process_axis:1)
    comm_spec = CommSpec(CollectiveCommPattern.SPLIT_FWD_GATHER_BWD, sharding_spec, shard_dim=1, logical_process_axis=1)
    tensor_to_shard = comm_spec.covert_spec_to_action(tensor_to_shard)
    if rank in (0, 2):
        assert tensor_to_shard.equal(sharded_tensor_to_comm_0)
    if rank in (1, 3):
        assert tensor_to_shard.equal(sharded_tensor_to_comm_1)