code: def all_reduce_mean(x: float, world_size: int) -> float:
      if world_size == 1:
          return x
      tensor = torch.tensor([x], device=get_current_device())
      dist.all_reduce(tensor)
      tensor = tensor / world_size
      return tensor.item()
  
file_path: ColossalAI/examples/language/llama2/performance_evaluator.py
imports: from time import time
  from typing import Optional
  import torch
  import torch.distributed
  from torch import Tensor
  import colossalai.utils.device
  from colossalai.cluster import DistCoordinator
  from colossalai.utils.device import get_current_device
label_desc: Use image node_id 2 for calling a global function with example usage: all_reduce_mean(x, world_size) and returns: tensor, x