    def forward(
        self,
        query_states: torch.Tensor,
        past_key_states: Optional[torch.Tensor] = None,
        past_value_states: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.BoolTensor] = None,
        source_states: Optional[torch.Tensor] = None,
        source_attention_mask: Optional[torch.BoolTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        position_bias: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        query_length: Optional[int] = None,
    ):
        """
        # Parameters

        query_states : `torch.Tensor`
            Shape `batch_size x seq_len x hidden_dim`
        past_key_states : `torch.Tensor`, optional
            Shape `batch_size x seq_len x hidden_dim`
            These are the key_states from the previous step of the decoder.
        past_value_states : `torch.Tensor`, optional
            Shape `batch_size x seq_len x hidden_dim`
            These are the value_states from the previous step of the decoder.
        attention_mask : `torch.BoolTensor`, optional
            Shape `batch_size x seq_len`
        source_states : `torch.Tensor`, optional
            Shape `batch_size x source_seq_len x hidden_dim`
            This is from the final state of attention over the source (encoder);
            it is passed when this module is being used for cross-attention.
        source_attention_mask : `torch.BoolTensor`, optional
            Shape `batch_size x source_seq_len`
        head_mask : `torch.BoolTensor`, optional
        position_bias : `torch.Tensor`, optional
        output_attentions : `bool`
            Whether to also return the attention probabilities, default = `False`

        !!! Note
            `source_states` needs to be passed in case of cross-attention.

        """
        query_layer = self._query_layer(query_states)

        key_layer = self._project(
            query_states,
            self.key,
            source_states,
            past_key_states,
        )

        value_layer = self._project(
            query_states,
            self.value,
            source_states,
            past_value_states,
        )

        if self.is_cross_attention:
            attention_mask = source_attention_mask

        seq_lengths = self._get_lengths(query_states, past_key_states, source_states, query_length)

        attention_probs, position_bias = self._get_attention_probs(
            query_layer,
            key_layer,
            attention_mask,
            head_mask,
            seq_lengths,
            position_bias,
            past_key_states,
        )

        context_layer = self._output_layer(attention_probs, value_layer)

        present_key_value_state = (
            (key_layer, value_layer) if (self.is_decoder and use_cache) else None
        )

        if not output_attentions:
            attention_probs = None

        outputs = AttentionOutput(
            context_layer, present_key_value_state, position_bias, attention_probs
        )

        return outputs
# Parameters

query_states : `torch.Tensor`
Shape `batch_size x seq_len x hidden_dim`
past_key_states : `torch.Tensor`, optional
Shape `batch_size x seq_len x hidden_dim`
These are the key_states from the previous step of the decoder.
past_value_states : `torch.Tensor`, optional
Shape `batch_size x seq_len x hidden_dim`
These are the value_states from the previous step of the decoder.
attention_mask : `torch.BoolTensor`, optional
Shape `batch_size x seq_len`
source_states : `torch.Tensor`, optional
Shape `batch_size x source_seq_len x hidden_dim`
This is from the final state of attention over the source (encoder);
it is passed when this module is being used for cross-attention.
source_attention_mask : `torch.BoolTensor`, optional
Shape `batch_size x source_seq_len`
head_mask : `torch.BoolTensor`, optional
position_bias : `torch.Tensor`, optional
output_attentions : `bool`
Whether to also return the attention probabilities, default = `False`

!!! Note
`source_states` needs to be passed in case of cross-attention.
allennlp/modules/transformer/attention_module.py