def forward(
        self,
        x,
        *,
        mask = None
    ):
        """
        Forward pass for the MossFormer module.

        Args:
            x (Tensor): Input tensor of shape (B, T, Q, C) where:
                B = batch size,
                T = total sequence length,
                Q = number of query features,
                C = feature dimension.
            mask (Tensor, optional): Attention mask for padding.

        Returns:
            Tensor: Output tensor of shape (B, T, C).
        """
        
        # Unpack input dimensions
        B, T, Q, C = x.size()
        x = x.view(B*T, Q, C)  # Reshape input for processing

        # Prenormalization step
        normed_x = x 

        # Optionally shift tokens for better performance
        residual = x  # Store residual for skip connection

        if self.shift_tokens:
            # Split and shift tokens for enhanced information flow
            x_shift, x_pass = normed_x.chunk(2, dim = -1)
            x_shift = F.pad(x_shift, (0, 0, 1, -1), value = 0.)  # Pad to maintain shape
            normed_x = torch.cat((x_shift, x_pass), dim = -1)

        # Initial projections to hidden space
        v, u = self.to_hidden(normed_x).chunk(2, dim = -1)  # Split into two tensors
        qk = self.to_qk(normed_x)  # Project to query/key dimensions

        # Offset and scale for attention
        quad_q, lin_q, quad_k, lin_k = self.qk_offset_scale(qk)
        att_v, att_u = self.cal_attention(x, quad_q, lin_q, quad_k, lin_k, v, u, B)

        # Gate the outputs and apply skip connection
        out = (att_u * v) * self.gateActivate(att_v * u)       
        x = x + self.to_out(out)  # Combine with residual
        return x

        Forward pass for the MossFormer module.

        Args:
            x (Tensor): Input tensor of shape (B, T, Q, C) where:
                B = batch size,
                T = total sequence length,
                Q = number of query features,
                C = feature dimension.
            mask (Tensor, optional): Attention mask for padding.

        Returns:
            Tensor: Output tensor of shape (B, T, C).
        
clearvoice/clearvoice/models/mossformer_gan_se/mossformer.py