def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Forward pass with residual connection and layer normalization.

        Args:
            hidden_states (torch.Tensor): Input tensor of shape [batch_size, seq_len, hidden_size]
            attention_mask (Optional[torch.Tensor]): Attention mask (currently not supported)

        Returns:
            torch.Tensor: Output tensor of shape [batch_size, seq_len, hidden_size]
        """
        normed_hidden_states = self.layer_norm(hidden_states)

        attn_output = self.attention(normed_hidden_states, attention_mask)

        if self.dropout is not None:
            attn_output = self.dropout(attn_output)

        output = hidden_states + attn_output

        return output

        Forward pass with residual connection and layer normalization.

        Args:
            hidden_states (torch.Tensor): Input tensor of shape [batch_size, seq_len, hidden_size]
            attention_mask (Optional[torch.Tensor]): Attention mask (currently not supported)

        Returns:
            torch.Tensor: Output tensor of shape [batch_size, seq_len, hidden_size]
        
src/liger_kernel/transformers/fused_neighborhood_attention.py