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

[((329, 395), 'fmoe.functions.count_by_gate', 'count_by_gate', (['topk_idx', 'num_expert', 'world_size'], {'require_pos': '(False)'}), '(topk_idx, num_expert, world_size, require_pos=False)\n', (342, 395), False, 'from fmoe.functions import count_by_gate\n'), ((423, 491), 'fmoe_cuda.limit_by_capacity', 'fmoe_native.limit_by_capacity', (['gec', 'capacity', 'num_expert', 'world_size'], {}), '(gec, capacity, num_expert, world_size)\n', (452, 491), True, 'import fmoe_cuda as fmoe_native\n'), ((219, 284), 'torch.ones', 'torch.ones', (['num_expert'], {'dtype': 'torch.int32', 'device': 'topk_idx.device'}), '(num_expert, dtype=torch.int32, device=topk_idx.device)\n', (229, 284), False, 'import torch\n'), ((546, 606), 'fmoe_cuda.expert_exchange', 'fmoe_native.expert_exchange', (['new_gec', 'num_expert', 'world_size'], {}), '(new_gec, num_expert, world_size)\n', (573, 606), True, 'import fmoe_cuda as fmoe_native\n')]