code: def _all_reduce(tensor, comm_spec, async_op=False):
      """
      Implement all reduce operation on device mesh based on information provided by comm_spec.
      """
      process_groups = (
          comm_spec.device_mesh.get_process_group_for_all_axes()
      )
      process_group = process_groups[comm_spec.logical_process_axis]
  
      if not tensor.is_contiguous():
          tensor = tensor.contiguous()
      dist.all_reduce(
          tensor,
          op=ReduceOp.SUM,
          group=process_group,
          async_op=async_op,
      )
      return tensor
  
imports: import operator
  from enum import Enum
  from functools import reduce
  import torch
  import torch.distributed
  from torch.distributed import ReduceOp
label_desc: Use image node_id 4 for calling a global function with example usage: _all_reduce(tensor, comm_spec, async_op) and returns: tensor