r"""
Utilities that may be used in the gates
"""
import torch
from fmoe.functions import count_by_gate
import fmoe_cuda as fmoe_native


def limit_by_capacity(topk_idx, num_expert, world_size, capacity):
    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

    fmoe_native.prune_gate_by_capacity(topk_idx,
            new_lec.to(torch.int32), num_expert, world_size)

    return new_lec, new_gec

[((13, 20, 14, 30), 'fmoe.functions.count_by_gate', 'count_by_gate', (), '', False, 'from fmoe.functions import count_by_gate\n'), ((15, 15, 16, 35), 'fmoe_cuda.limit_by_capacity', 'fmoe_native.limit_by_capacity', ({(15, 45, 15, 48): 'gec', (15, 50, 15, 58): 'capacity', (16, 12, 16, 22): 'num_expert', (16, 24, 16, 34): 'world_size'}, {}), '(gec, capacity, num_expert, world_size)', True, 'import fmoe_cuda as fmoe_native\n'), ((10, 15, 11, 35), 'torch.ones', 'torch.ones', (), '', False, 'import torch\n'), ((18, 19, 18, 79), 'fmoe_cuda.expert_exchange', 'fmoe_native.expert_exchange', ({(18, 47, 18, 54): 'new_gec', (18, 56, 18, 66): 'num_expert', (18, 68, 18, 78): 'world_size'}, {}), '(new_gec, num_expert, world_size)', True, 'import fmoe_cuda as fmoe_native\n')]