code: def _gather_along_last_dim(input_):
      """Gather tensors and concatinate along the last dimension."""
  
      world_size = get_tensor_model_parallel_world_size()
      # Bypass the function if we are using only 1 GPU.
      if world_size == 1:
          return input_
  
      # Size and dimension.
      last_dim = input_.dim() - 1
      rank = get_tensor_model_parallel_rank()
  
      tensor_list = [
          torch.empty_like(input_) for _ in range(world_size)
      ]
      tensor_list[rank] = input_
      torch.distributed.all_gather(
          tensor_list, input_, group=get_tensor_model_parallel_group()
      )
  
      # Note: torch.cat already creates a contiguous tensor.
      output = torch.cat(tensor_list, dim=last_dim).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 4 for calling a global function with example usage: _gather_along_last_dim(input_) and returns: output, input_