    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
    ) -> TransformerStackOutput:
        """
        # Parameters

        hidden_states : `torch.Tensor`
            Shape `batch_size x seq_len x hidden_dim`
        attention_mask : `torch.BoolTensor`, optional
            Shape `batch_size x seq_len`
        head_mask : `torch.BoolTensor`, optional
        output_attentions : `bool`
            Whether to also return the attention probabilities, default = `False`
        output_hidden_states : `bool`
            Whether to return the hidden_states for all layers, default = `False`
        """
        all_hidden_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self._add_cross_attention else None
        for i, layer_module in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)  # type: ignore

            layer_head_mask = head_mask[i] if head_mask is not None else None

            layer_outputs = layer_module(
                hidden_states,
                attention_mask,
                layer_head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                output_attentions,
            )
            hidden_states = layer_outputs.hidden_states
            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)  # type: ignore
                if self._add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)  # type: ignore

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)  # type: ignore

        return TransformerStackOutput(
            hidden_states, all_hidden_states, all_attentions, all_cross_attentions
        )
# Parameters

hidden_states : `torch.Tensor`
Shape `batch_size x seq_len x hidden_dim`
attention_mask : `torch.BoolTensor`, optional
Shape `batch_size x seq_len`
head_mask : `torch.BoolTensor`, optional
output_attentions : `bool`
Whether to also return the attention probabilities, default = `False`
output_hidden_states : `bool`
Whether to return the hidden_states for all layers, default = `False`
allennlp/modules/transformer/transformer_stack.py