code: def count_by_gate(gate, num_expert, world_size, require_pos=True):
      with torch.no_grad():
          local_expert_count = torch.zeros(
              num_expert * world_size,
              device=gate.device,
              dtype=torch.int32,
          )
          fmoe_cuda.expert_count(gate, local_expert_count)
          local_expert_count = local_expert_count.long()
  
          if world_size > 1:
              global_expert_count = fmoe_cuda.expert_exchange(
                  local_expert_count, num_expert, world_size
              )
          else:
              global_expert_count = local_expert_count
          if not require_pos:
              pos = None
          else:
              lec_cum = torch.cumsum(local_expert_count, dim=0).int()
              pos_size = lec_cum[-1].item()
              pos = torch.empty(
                  (pos_size,), device=gate.device, dtype=torch.long
              )
              fmoe_cuda.assign_pos(lec_cum, gate, pos)
      return pos, local_expert_count, global_expert_count
  
imports: import torch
  from torch.autograd import Function
  import fmoe_cuda
  from .utils import get_torch_default_comm
label_desc: Use image node_id 3 for calling a global function with example usage: count_by_gate(gate, num_expert, world_size, require_pos) and returns: pos,  local_expert_count,  global_expert_count