    def forward(  # type: ignore
        self,
        hidden_states: torch.Tensor,
        mask: Optional[torch.BoolTensor] = None,
        key_value_states: Optional[FloatT] = None,
        position_bias: Optional[FloatT] = None,
        past_key_value: Optional[
            Tuple[FloatT, FloatT]
        ] = None,  # this is used when taking decoding steps.
        layer_head_mask: Optional[BoolT] = None,
        query_length: Optional[int] = None,  # only relevant in cross-attention.
        use_cache: bool = False,
        output_attentions: bool = False,
    ) -> AttentionOutput:
        """
        Self-attention (if key_value_states is None) or attention over source sentence (provided by
        key_value_states).
        """
        if past_key_value:
            past_key_states = past_key_value[0]
            past_value_states = past_key_value[1]
        else:
            past_key_states = None
            past_value_states = None

        outputs = super().forward(
            query_states=hidden_states,
            past_key_states=past_key_states,
            past_value_states=past_value_states,
            attention_mask=mask,
            source_states=key_value_states,
            source_attention_mask=None,  # TODO: is this a bug in current T5 code?
            head_mask=layer_head_mask,
            position_bias=position_bias,
            output_attentions=output_attentions,
            use_cache=use_cache,
            query_length=query_length,
        )

        return outputs
Self-attention (if key_value_states is None) or attention over source sentence (provided by
key_value_states).
allennlp/modules/transformer/attention_module.py