code: def gather_scalar_from_all_ranks(
      fn, world_size=1
  ):  # pylint: disable=invalid-name
      """
      Run a callable 'fn' on all ranks and gather the results on the specified rank.
  
      Args:
      - fn (callable): A function that computes the value. This should not have any side effects.
      - rank (int, optional): The rank that gathers the values. Default is 0.
      - world_size (int, optional): Total number of processes in the current distributed setup.
  
      Returns:
      - A list of computed values from all ranks if on the gathering rank, otherwise None.
      """
      value_scalar = fn()
      if not is_distributed():
          return [value_scalar]
      value_tensor = torch.tensor(
          value_scalar, device=torch.cuda.current_device()
      ).float()
  
      if not is_main_process():
          dist.gather(value_tensor, dst=0)
      else:
          gathered_tensors = [
              torch.zeros_like(value_tensor) for _ in range(world_size)
          ]
          dist.gather(value_tensor, gather_list=gathered_tensors, dst=0)
  
          # Convert tensors back to their original type (int or float)
          gathered_values = []
          for tensor in gathered_tensors:
              if tensor == tensor.int():
                  gathered_values.append(int(tensor.item()))
              else:
                  gathered_values.append(float(tensor.item()))
          return gathered_values
      return None
  
file_path: axolotl/src/axolotl/utils/distributed.py
imports: import os
  import pickle
  from contextlib import contextmanager
  import torch
  import torch.distributed
  from accelerate import Accelerator
label_desc: Use image node_id 8 for calling a global function with example usage: gather_scalar_from_all_ranks(fn, world_size) and returns: None, list, gathered_values