code: def assert_equal_in_group(
      tensor: Tensor, process_group: ProcessGroup = None
  ):
      # all gather tensors from different ranks
      world_size = dist.get_world_size(process_group)
      tensor_list = [
          torch.empty_like(tensor) for _ in range(world_size)
      ]
      dist.all_gather(tensor_list, tensor, group=process_group)
  
      # check if they are equal one by one
      for i in range(world_size - 1):
          a = tensor_list[i]
          b = tensor_list[i + 1]
          assert torch.all(
              a == b
          ), f"expected tensors on rank {i} and {i + 1} to be equal but they are not, {a} vs {b}"
  
file_path: ColossalAI/colossalai/testing/comparison.py
imports: from typing import Any, List, OrderedDict
  import torch
  import torch.distributed
  from torch import Tensor
  from torch.distributed import ProcessGroup
  from torch.testing import assert_close
  from torch.utils._pytree import tree_flatten
label_desc: Use image node_id 4 for calling a global function with example usage: assert_equal_in_group(tensor, process_group) without return types