code: def _split(tensor, comm_spec):
      """
      Implement shard 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]
  
      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).contiguous()
      return output
  
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 2 for calling a global function with example usage: _split(tensor, comm_spec) and returns: output