import math
import numpy
import megengine as mge
import megengine.module as M
import megengine.functional as F

from .layer_norm import LayerNorm


class DecoderLayer(M.Module):
    """Single decoder layer module."""

    def __init__(
        self,
        size,
        self_attn,
        src_attn,
        feed_forward,
        dropout_rate,
        normalize_before=True,
    ):
        """Construct an DecoderLayer object."""
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.norm1 = LayerNorm(size)
        self.norm2 = LayerNorm(size)
        self.norm3 = LayerNorm(size)
        self.dropout = 
M.dropout.Dropout(dropout_rate)