import megengine as mge
import megengine.module as M
from megengine import functional as F
import numpy as np
from .transformer import MultiheadAttention
#from .utility import has_nan_or_inf

# mge.core.set_option('async_level', 0)

class DecoderWrapper(M.Module):
    def __init__(self, cfg):
        super().__init__()
        channels = cfg.distiller.HIDDEN_DIM
        heads = cfg.distiller.ATT_HEADS

        # this is a local module derived from official implementation, we modify the last modules
        self.matt = MultiheadAttention(channels, heads)

        self.pos_projector = 
M.Linear(in_features=channels, out_features=channels)