def forward(
        self,
        x,
        state=None, # TODO consolidate with incremental_state
        padding_mask: Optional[torch.Tensor] = None,  # Mega: encoder_padding_mask
        need_weights: bool = False,
        attn_mask: Optional[torch.Tensor] = None,
        before_attn_fn: bool = False,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """Input shape: (B L D) = batch, length, dimension

        Dimensions:
            B: Batch size
            L: Sequence length (l_seq)
            C: Chunk size
            K: Number of chunks (L / C)
            D: Model dimension (d_model)     ($d$ in paper)
            V: Dim. attention output (Mega: paper $v$, code `d_attout`, annotation E)
            Z: Dim. attention input  (Mega: paper $z$, code `d_attin`, annotation S)
        """
        if self.transposed: x = x.transpose(-1, -2)  # (B L D)
        B, L, D = x.size()
        assert D == self.d_model

        residual = x
        if self.prenorm:
            x = self.norm(x)

        v = self.activation(self.v_proj(x))  # (B L V)

        mx, _ = self.ssm(x, state=state, padding_mask=padding_mask)  # (B L D)
        # Original Mega code bakes a SiLU activation at the end of the MultiHeadEMA module
        # It has been moved here, which makes more sense to keep the SSM module linear
        # and so the activation is configurable and consistent with the rest of the block
        mx = self.activation(mx)
        mx = self.dropout(mx)

        base = self.mx_proj(mx)  # (B L D) -> (B L 2*D+Z+V)
        u, zr, hx = torch.split(base, [D, self.d_attin + self.d_attout, D], dim=-1)
        u = torch.sigmoid(u)  # (B L D) \phi (eq. 13)
        # Mega specifies to hard-code SiLU here, but the self.activation is always silu
        # in their configs anyways so this seems more sensible in case it's changed
        z, r = torch.split(self.activation(zr), [
            self.d_attin,  # z = (B L Z) Z (eq. 7)
            self.d_attout, # r = (B L V) \gamma (eq. 12)
        ], dim=-1)
        z = z.unsqueeze(2) * self.gamma + self.beta
        q, k = torch.unbind(z, dim=2)  # (B L Z) Q and K (eq. 8 and 9)

        q = q.unsqueeze(1)  # (B 1 L Z)
        k = k.unsqueeze(1)  # (B 1 L Z)
        v = v.unsqueeze(1)  # (B 1 L Z)
        if self.l_chunk < 0:
            if padding_mask is not None:
                padding_mask = padding_mask.unsqueeze(1)  # (B 1 L)
        else:
            if L < self.l_chunk:
                pass
            else:
                q = rearrange(q, 'b 1 (k c) z -> b k c z', c=self.l_chunk)

            l_ctx = k.size(2)  # Transcribed from orig, why is this not the same as L?
            if l_ctx < self.l_chunk:
                if padding_mask is not None:
                    padding_mask = padding_mask.unsqueeze(1)  # (B 1 C)?
            else:
                k = rearrange(k, 'b 1 (k c) z -> b k c z', c=self.l_chunk)
                v = rearrange(v, 'b 1 (k c) z -> b k c z', c=self.l_chunk)
                if padding_mask is not None:
                    padding_mask = rearrange(padding_mask, 'b (k c) -> b k c', c=self.l_chunk)

        # This is part of a workaround to get around fork/join parallelism not supporting Optional types.
        if padding_mask is not None and padding_mask.dim() == 0:
            padding_mask = None

        if self.attention_activation_fn is None:  # Softmax case
            attn_weights = self.softmax_attention(q, k, padding_mask, attn_mask, before_attn_fn)
        else:
            attn_weights = self.element_attention(q, k, padding_mask, attn_mask, before_attn_fn)

        if before_attn_fn:
            if self.transposed: v = v.transpose(-1, -2)
            # return attn_weights, v  # TODO looks like bug in orig code
            return v, attn_weights

        v = self.drop_attout(v)  # (B K C V)
        kernel = self.drop_attin(attn_weights)  # (B K C C)
        h = rearrange(torch.matmul(kernel, v), 'b k c v -> b (k c) v') # (B L V)
        h = self.activation(hx + self.h_proj(h * r))  # (B L D)
        h = self.dropout(h)
        # Output (y) from update gate u (\phi): u * h + (1-u) * x, eq. (15)
        out = torch.addcmul(residual, u, h - residual)  # (B L D)

        if not self.prenorm:
            out = self.norm(out)

        if self.transposed: out = out.transpose(-1, -2)

        # FFN
        out, _ = self.nffn(out, state=None)


        if not need_weights: attn_weights = None
        # Because this class expects to return a state, it's a little inconvenient to return attention weights.
        # The official Mega code doesn't return it either.
        return out, _ # , attn_weights
Input shape: (B L D) = batch, length, dimension

        Dimensions:
            B: Batch size
            L: Sequence length (l_seq)
            C: Chunk size
            K: Number of chunks (L / C)
            D: Model dimension (d_model)     ($d$ in paper)
            V: Dim. attention output (Mega: paper $v$, code `d_attout`, annotation E)
            Z: Dim. attention input  (Mega: paper $z$, code `d_attin`, annotation S)
        
analyze/convnexts4nd/src/models/sequence/modules/megablock.py