code: def _split_along_first_dim(input_):
      """Split the tensor along its first dimension and keep the
      corresponding slice."""
  
      world_size = get_tensor_model_parallel_world_size()
      # Bypass the function if we are using only 1 GPU.
      if world_size == 1:
          return input_
  
      # Split along first dimension.
      dim_size = input_.size()[0]
      assert (
          dim_size % world_size == 0
      ), "First dimension of the tensor should be divisible by tensor parallel size"
      local_dim_size = dim_size // world_size
      rank = get_tensor_model_parallel_rank()
      dim_offset = rank * local_dim_size
  
      output = input_[
          dim_offset : dim_offset + local_dim_size
      ].contiguous()
  
      return output
  
imports: import torch
  from megatron.core.parallel_state import get_tensor_and_expert_parallel_group, get_tensor_model_parallel_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
  from .utils import split_tensor_along_last_dim
label_desc: Use image node_id 3 for calling a global function with example usage: _split_along_first_dim(input_) and returns: output, input_