    def forward(
            self,
            data: torch.Tensor
    ) -> torch.Tensor:
        """
        Forward pass.
        Args:
            data (`torch.Tensor`):
                Input tensor of the shape [B_, N, C].
        Returns:
            output (`torch.Tensor`):
                Output tensor of the shape [B_, N, C].
        """
        # Get shape of input
        b_, n, _ = data.shape
        # Perform query key value mapping
        qkv = self.qkv_mapping(data).reshape(b_, n, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        # Scale query
        q = q * self.scale
        # Compute attention maps
        attn = self.softmax(q @ k.transpose(-2, -1) + self._get_relative_positional_bias())
        # Map value with attention maps
        output = (attn @ v).transpose(1, 2).reshape(b_, n, -1)
        # Perform final projection and dropout
        output = self.proj(output)
        output = self.proj_drop(output)
        return output
Forward pass.
Args:
data (`torch.Tensor`):
Input tensor of the shape [B_, N, C].
Returns:
output (`torch.Tensor`):
Output tensor of the shape [B_, N, C].
towhee/models/layers/relative_self_attention.py