def forward(
        self,
        u,
        state: Optional[Dict[str, Dict[str, Optional[torch.Tensor]]]] = None,
        padding_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, None]:
        """Input shape: Time x Batch x Channel
        Args:
            padding_mask (ByteTensor, optional): mask to exclude
                keys that are pads, of shape `(batch, src_len)`, where
                padding elements are indicated by 1s.
        """

        B, L, H = u.size()
        assert H == self.H

        u = u.transpose(-1, -2)  # (B H L)
        if padding_mask is not None:
            u = u * (1.0 - padding_mask.unsqueeze(1).type_as(u))

        # assert not self.bidirectional or state is None, 'Bidirectional EMA does not support incremental state'
        if state is not None:
            raise NotImplementedError(
                "MultiHeadEMA module does not support state passing in this repository."
                "Use S4D for more functionality such as state passing."
            )
        else:
            k = self.kernel(L)  # (H L)
            l_fft = L
            s = 0
            l_kernel = k.size(1)
            assert l_kernel == L
            u_ = u
            if self.bidirectional:
                # This is twice as inefficient as it could be
                # See S4 FFT conv bidirectional implementation for improvement
                k1, k2 = torch.split(k, [self.H, self.H], dim=0)
                k = F.pad(k1, (l_kernel - 1, 0)) + F.pad(k2.flip(-1), (0, l_kernel - 1))  # (H 2*L-1)
                u_ = F.pad(u, (l_kernel - 1, 0))
                l_fft = l_fft + l_kernel - 1
                s = 2 * l_kernel - 2

            k_f = torch.fft.rfft(k.float(), n=2 * l_fft)
            u_f = torch.fft.rfft(u_.float(), n=2 * l_fft)
            y = torch.fft.irfft(u_f * k_f, n=2 * l_fft)[..., s:s + L]  # (B H L)
            y = y.type_as(u)
            y = y + u * self.omega.unsqueeze(-1)  # (B H L)
            y = y.transpose(-1, -2)

        return y, None  # empty state
Input shape: Time x Batch x Channel
        Args:
            padding_mask (ByteTensor, optional): mask to exclude
                keys that are pads, of shape `(batch, src_len)`, where
                padding elements are indicated by 1s.
        
analyze/convnexts4nd/src/models/sequence/modules/megablock.py