class_docstr: Basic MoE layer.
class_docstr_tok: ["Basic","MoE","layer","."]
docstr: Gather tensors and concatenate along the first dimension.
docstr_tok: ["Gather","tensors","and","concatenate","along","the","first","dimension","."]
code: def gather_indices(self, local_indices):
      """Gather tensors and concatenate along the first dimension."""
      group = get_tensor_and_expert_parallel_group()
      world_size = torch.distributed.get_world_size(group=group)
      # Bypass the function if we are using only 1 GPU.
      if world_size == 1:
          return local_indices
  
      dim_size = list(local_indices.size())
      dim_size[0] = dim_size[0] * world_size
  
      # TODO pre allocate memory
      output = torch.empty(
          dim_size,
          dtype=local_indices.dtype,
          device=torch.cuda.current_device(),
      )
      torch.distributed._all_gather_base(
          output, local_indices.contiguous(), group=group
      )
      return output
  
file_path: megatron-lm/megatron/core/transformer/moe/base_moe_layer.py
imports: from abc import ABC, abstractmethod
  import torch
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.parallel_state import get_tensor_and_expert_parallel_group
  from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_data_parallel_rng_tracker_name
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.transformer_config import TransformerConfig
label_desc: Use image node_id 2 for calling the BaseMoELayer obj's underlying member method code with example usage: obj.gather_indices(local_indices) and returns: output, local_indices