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, 'fmoe.functions.count_by_gate', 'count_by_gate', False, 'from fmoe.functions import count_by_gate\n'), (15, 'fmoe_cuda.limit_by_capacity', 'fmoe_native.limit_by_capacity', True, 'import fmoe_cuda as fmoe_native\n'), (10, 'torch.ones', 'torch.ones', False, 'import torch\n'), (18, 'fmoe_cuda.expert_exchange', 'fmoe_native.expert_exchange', True, 'import fmoe_cuda as fmoe_native\n')]