code: def _all_to_all(tensor: torch.Tensor, comm_spec: CommSpec):
      """
      Implement all to all operation on device mesh based on information provided by comm_spec.
      """
      process_group = comm_spec.process_group_dict[
          comm_spec.logical_process_axis
      ]
      world_size = dist.get_world_size(process_group)
      new_shape = list(tensor.shape)
      new_shape[comm_spec.shard_dim] = (
          new_shape[comm_spec.shard_dim] // world_size
      )
      new_shape = torch.Size(new_shape)
      output_tensor_list = [
          torch.zeros(
              new_shape, dtype=tensor.dtype, device=tensor.device
          )
          for _ in range(world_size)
      ]
      dim = comm_spec.shard_dim
      length = tensor.shape[comm_spec.shard_dim] // world_size
      input_tensor_list = [
          torch.narrow(tensor, dim, length * i, length).contiguous()
          for i in range(world_size)
      ]
      group = process_group
      dist.all_to_all(output_tensor_list, input_tensor_list, group)
      output = torch.cat(
          tuple(output_tensor_list), comm_spec.gather_dim
      ).contiguous()
      return output
  
imports: from enum import Enum
  from typing import Dict
  import torch
  import torch.distributed
  from torch.distributed import ReduceOp
label_desc: Use image node_id 3 for calling a global function with example usage: _all_to_all(tensor, comm_spec) and returns: output