code: def limit_by_capacity(topk_idx, num_expert, world_size, capacity):
      with torch.no_grad():
          capacity = (
              torch.ones(
                  num_expert, dtype=torch.int32, device=topk_idx.device
              )
              * capacity
          )
  
          pos, lec, gec = count_by_gate(
              topk_idx, num_expert, world_size, require_pos=False
          )
          new_gec = fmoe_native.limit_by_capacity(
              gec, capacity, num_expert, world_size
          )
          if world_size > 1:
              new_lec = fmoe_native.expert_exchange(
                  new_gec, num_expert, world_size
              )
          else:
              new_lec = new_gec
  
          topk_idx = fmoe_native.prune_gate_by_capacity(
              topk_idx, new_lec.to(torch.int32), num_expert, world_size
          )
      return new_lec, new_gec, topk_idx
  
imports: import torch
  from fmoe.functions import count_by_gate
  import fmoe_cuda
label_desc: Use image node_id 1 for calling a global function with example usage: limit_by_capacity(topk_idx, num_expert, world_size, capacity) and returns: new_lec,  new_gec,  topk_idx