def _T5Attention_forward(
    self,
    hidden_states,
    mask=None,
    key_value_states=None,
    position_bias=None,
    past_key_value=None,
    layer_head_mask=None,
    query_length=None,
    use_cache=True,
    output_attentions=False,
):
    """
    Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
    """
    # Input is (batch_size, seq_length, dim)
    # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length)
    # past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head)
    batch_size, seq_length = hidden_states.shape[:2]

    def shape(states):
        """projection"""
        return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim)

    def project(hidden_states, proj_layer, key_value_states, past_key_value):
        """projects hidden states correctly to key/query states"""
        if key_value_states is None:
            # self-attn
            hidden_states = shape(proj_layer(hidden_states))
        else:
            if past_key_value is None:
                hidden_states = shape(proj_layer(key_value_states))
            else:
                hidden_states = past_key_value.transpose(0, 1)[
                    :, : key_value_states.shape[1], :, :
                ]
        return hidden_states

    # get query states
    query = shape(self.q(hidden_states))

    # get key/value states
    key = project(
        hidden_states,
        self.k,
        key_value_states,
        past_key_value[1] if past_key_value is not None else None,
    ).to(dtype=query.dtype)
    value = project(
        hidden_states,
        self.v,
        key_value_states,
        past_key_value[2] if past_key_value is not None else None,
    ).to(dtype=query.dtype)

    real_seq_length = seq_length
    if past_key_value is not None:
        real_seq_length += (
            past_key_value[0].shape[2] if query_length is None else query_length
        )

    key_length = (
        real_seq_length if key_value_states is None else key_value_states.shape[1]
    )
    if position_bias is None:
        if not self.has_relative_attention_bias:
            position_bias = torch.zeros(
                (1, self.n_heads, real_seq_length, key_length),
                device=hidden_states.device,
                dtype=hidden_states.dtype,
            )
            if self.gradient_checkpointing and self.training:
                position_bias.requires_grad = True
        else:
            position_bias = self.compute_bias(
                real_seq_length, key_length, device=hidden_states.device
            )

        # if key and values are already calculated
        # we want only the last query position bias
        if past_key_value is not None:
            position_bias = position_bias[:, :, -hidden_states.size(1) :, :]

        if mask is not None:
            position_bias = torch.tensor(position_bias) + torch.tensor(
                mask
            )  # (batch_size, n_heads, seq_length, key_length)

    if self.pruned_heads:
        mask = torch.ones(position_bias.shape[1])
        mask[list(self.pruned_heads)] = 0
        position_bias_masked = position_bias[:, mask.bool()]
    else:
        position_bias_masked = position_bias

    if key_value_states is None:
        decoded_tokens = (
            torch.tensor(past_key_value[0].size(-2))
            if past_key_value is not None
            else None
        )
    else:
        decoded_tokens = torch.zeros(1, dtype=torch.long).contiguous()[0]
    (
        attn_output,
        attn_weights,
        present_key_value_state,
    ) = self._IPEXScaleDotProduct(
        query,
        key,
        value,
        1,
        past_key_value,
        layer_head_mask,
        position_bias_masked,
        None,
        False,
        decoded_tokens,
    )
    if not (self.is_decoder and use_cache):
        present_key_value_state = None

    attn_output = attn_output.transpose(1, 2)
    attn_output = attn_output.reshape(query.shape[0], query.shape[1], -1)
    attn_output = self.o(attn_output)

    outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)

    if output_attentions:
        outputs = outputs + (attn_weights,)
    return outputs

    Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
    
intel_extension_for_pytorch/transformers/models/reference/modules/attentions.py