code: def prepare_forward(gate, num_expert, world_size):
      r"""
      Prepare necessary information from gate output for MoE computation.
  
      Args:
          gate: a 1-d Long Tensor representing the target expert of each input
          sample.
          num_expert: number of experts on each worker.
          world_size: number of workers that hold different experts.
          comm: the communicator of all workers in the expert-parallel group.
      """
      pos, local_expert_count, global_expert_count = count_by_gate(
          gate, num_expert, world_size
      )
      with torch.no_grad():
          fwd_expert_count = global_expert_count.view(
              world_size, num_expert
          ).sum(dim=0)
          fwd_batch_size = int(fwd_expert_count.sum().item())
      return (
          pos,
          local_expert_count.cpu(),
          global_expert_count.cpu(),
          fwd_expert_count.cpu(),
          fwd_batch_size,
      )
  
imports: import torch
  from torch.autograd import Function
  import fmoe_cuda
  from .utils import get_torch_default_comm
label_desc: Use image node_id 4 for calling a global function with example usage: prepare_forward(gate, num_expert, world_size) and returns: pos,  local_expert_count,  global_expert_count,  fwd_expert_count,  fwd_batch_size