def forward(self, x, *, mask=None):
        """
        Forward pass for FLASH layer.
        
        Args:
            x (Tensor): Input tensor of shape (batch, seq_len, features).
            mask (Tensor, optional): Mask for attention. Defaults to None.
        
        Returns:
            Tensor: Output tensor after applying attention and projections.
        """
        
        # Pre-normalization step
        normed_x = x 
        residual = x  # Save residual for skip connection

        # Token shifting if enabled
        if self.shift_tokens:
            x_shift, x_pass = normed_x.chunk(2, dim=-1)
            x_shift = F.pad(x_shift, (0, 0, 1, -1), value=0.)
            normed_x = torch.cat((x_shift, x_pass), dim=-1)

        # Initial projections
        v, u = self.to_hidden(normed_x).chunk(2, dim=-1)
        qk = self.to_qk(normed_x)

        # Offset and scale
        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)

        # Output calculation with gating
        out = (att_u * v) * self.gateActivate(att_v * u)       
        x = x + self.to_out(out)  # Residual connection
        return x

        Forward pass for FLASH layer.
        
        Args:
            x (Tensor): Input tensor of shape (batch, seq_len, features).
            mask (Tensor, optional): Mask for attention. Defaults to None.
        
        Returns:
            Tensor: Output tensor after applying attention and projections.
        
clearvoice/clearvoice/models/mossformer2_sr/mossformer2_block.py