code: def _split(tensor: torch.Tensor, comm_spec: CommSpec):
      """
      Implement shard operation on device mesh based on information provided by comm_spec.
      """
      process_group = comm_spec.process_group_dict[
          comm_spec.logical_process_axis
      ]
      dim = comm_spec.shard_dim
      length = tensor.shape[comm_spec.shard_dim] // dist.get_world_size(
          process_group
      )
      start = length * dist.get_rank(process_group)
      output = (
          torch.narrow(tensor, dim, start, length).clone().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 2 for calling a global function with example usage: _split(tensor, comm_spec) and returns: output