utils.py

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.reshape(-1), num_expert, world_size)
    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