src/transformers/models/yoso/modeling_yoso.py/YosoSelfAttention/forward
class YosoSelfAttention:    def forward(self, hidden_states, attention_mask=None, output_attentions=False):
        mixed_query_layer = self.query(hidden_states)

        key_layer = self.transpose_for_scores(self.key(hidden_states))
        value_layer = self.transpose_for_scores(self.value(hidden_states))
        query_layer = self.transpose_for_scores(mixed_query_layer)

        if self.use_conv:
            conv_value_layer = self.conv(value_layer * attention_mask[:, None, :, None])

        batch_size, num_heads, seq_len, head_dim = query_layer.size()

        query_layer = query_layer.reshape(batch_size * num_heads, seq_len, head_dim)
        key_layer = key_layer.reshape(batch_size * num_heads, seq_len, head_dim)
        value_layer = value_layer.reshape(batch_size * num_heads, seq_len, head_dim)

        # revert changes made by get_extended_attention_mask
        attention_mask = 1.0 + attention_mask / 10000.0
        attention_mask = (
            attention_mask.squeeze().repeat(1, num_heads, 1).reshape(batch_size * num_heads, seq_len).int()
        )

        # The CUDA kernels are most efficient with inputs whose size is a multiple of a GPU's warp size (32). Inputs
        # smaller than this are padded with zeros.
        gpu_warp_size = 32

        if (not self.use_expectation) and head_dim < gpu_warp_size:
            pad_size = batch_size * num_heads, seq_len, gpu_warp_size - head_dim

            query_layer = torch.cat(
                [
                    query_layer,
                    torch.zeros(pad_size, device=query_layer.device),
                ],
                dim=-1,
            )
            key_layer = torch.cat(
                [
                    key_layer,
                    torch.zeros(pad_size, device=key_layer.device),
                ],
                dim=-1,
            )
            value_layer = torch.cat(
                [
                    value_layer,
                    torch.zeros(pad_size, device=value_layer.device),
                ],
                dim=-1,
            )

        if self.use_expectation or self.training:
            query_layer, key_layer = normalize([query_layer, key_layer])

        if self.use_expectation:
            context_layer = YosoCumulation.apply(
                attention_mask, attention_mask, query_layer, key_layer, value_layer, self.lsh_config
            )
        else:
            context_layer = YosoLSHCumulation.apply(
                attention_mask, attention_mask, query_layer, key_layer, value_layer, self.lsh_config
            )

        if (not self.use_expectation) and head_dim < gpu_warp_size:
            context_layer = context_layer[:, :, :head_dim]

        context_layer = normalize(context_layer)

        context_layer = context_layer.reshape(batch_size, num_heads, seq_len, head_dim)

        if self.use_conv:
            context_layer += conv_value_layer

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (context_layer, context_layer) if output_attentions else (context_layer,)

        return outputs