Repository: pytorch
Issue: Improved Transformer and MultiHeadAttention design
Labels: module: nn,triaged,enhancement

## 🚀 Feature

Improved Transformer and MultiHeadAttention design

## Motivation

Current Transformer and MultiHeadAttention design is not a very `PyTorch`-ish. The `in_proj` is not even a `nn.Module`. `multihead_attention_forward` does not follow PEP-8.

## Pitch

Re-write Transformer and MultiHeadAttention

## Alternatives

No

## Additional context

```python
import math
import warnings

import torch
import torch.nn as nn
import torch.nn.functional as F

from torch import Tensor
from typing import Any, Callable, Dict, List, Optional, Tuple, Type


class MultiheadAttention(nn.Module):
    r"""Allows the model to jointly attend to information
    from different representation subspaces.
    See `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_
    .. math::
        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
    where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
    Args:
        embed_dim: total dimension of the model.
        num_heads: parallel attention heads.
        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
        bias: add bias as module parameter. Default: True.
        add_bias_kv: add bias to the key and value sequences at dim=0.
        add_zero_attn: add a new batch of zeros to the key and
                       value sequences at dim=1.
        k_dim: total number of features in key. Default: None.
        v_dim: total number of features in value. Default: None.
        batch_first: If ``True``, then the input and output tensors are provided
            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
    Note that if :attr:`k_dim` and :attr:`v_dim` are None, they will be set
    to :attr:`embed_dim` such that query, key, and value have the same
    number of features.
    Examples::
        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
    """
    __constants__ = ['batch_first']
    bias_k: Optional[torch.Tensor]
    bias_v: Optional[torch.Tensor]

    def __init__(self, embed_dim: int, num_heads: int,
                 attention_dropout: Optional[float] = 0.,
                 scale_factor: Optional[float] = 1.,
                 bias: Optional[bool] = True,
                 add_bias_kv: Optional[bool] = False,
                 add_zero_attn: Optional[bool] = False,
                 k_dim: Optional[int] = None, v_dim: Optional[int] = None,
                 batch_first: Optional[bool] = False,
                 **kwargs: Dict[str, Any]) -> None:
        super(MultiheadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.k_dim = k_dim if k_dim is not None else self.embed_dim
        self.v_dim = v_dim if v_dim is not None else self.embed_dim
        self._qkv_same_embed_dim = self.embed_dim == self.k_dim == self.v_dim

        self.num_heads = num_heads
        self.scale_factor = scale_factor
        self.batch_first = batch_first
        self.head_dim = self.embed_dim // self.num_heads
        self.scaling = float(self.head_dim * self.scale_factor) ** -0.5
        if not self.head_dim * self.num_heads == self.embed_dim:
            raise ValueError(f"embed_dim {self.embed_dim} not divisible by num_heads {self.num_heads}")

        self.in_proj = nn.Linear(self.embed_dim, self.embed_dim + self.k_dim + self.v_dim, bias=bias)
        self.dropout = nn.Dropout(attention_dropout)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=bias)

        if add_bias_kv:
            self.bias_k = nn.Parameter(torch.empty((1, 1, self.embed_dim)))
            self.bias_v = nn.Parameter(torch.empty((1, 1, self.embed_dim)))
        else:
            self.bias_k = self.bias_v = None

        self.add_zero_attn = add_zero_attn

        self._reset_parameters()

    def _reset_parameters(self):
        nn.init.xavier_uniform_(self.in_proj.weight)

        if self.in_proj.bias is not None:
            nn.init.constant_(self.in_proj.bias, 0.)
            nn.init.constant_(self.out_proj.bias, 0.)
        if self.bias_k is not None:
            nn.init.xavier_normal_(self.bias_k)
        if self.bias_v is not None:
            nn.init.xavier_normal_(self.bias_v)

    def __setstate__(self, state):
        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
        if '_qkv_same_embed_dim' not in state:
            state['_qkv_same_embed_dim'] = True

        super(MultiheadAttention, self).__setstate__(state)

    def forward(self, query: Tensor, key: Tensor, value: Tensor,
                attn_bias: Optional[Tensor] = None,
                attn_mask: Optional[Tensor] = None,
                key_padding_mask: Optional[Tensor] = None,
                need_weights: Optional[bool] = True,
                static_k: Optional[Tensor] = None,
                static_v: Optional[Tensor] = None,) -> Tuple[Tensor, Optional[Tensor]]:
        r"""
    Args:
        query, key, value: map a query and a set of key-value pairs to an output.
            See "Attention Is All You Need" for more details.
        attn_bias: 2D or 3D mask that add bias to attention output weights. Used for relative positional embedding.
            A 2D bias will be broadcasted for all the batches while a 3D mask allows to specify a different mask for
            the entries of each batch.
        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
        key_padding_mask: if provided, specified padding elements in the key will
            be ignored by the attention. When given a binary mask and a value is True,
            the corresponding value on the attention layer will be ignored. When given
            a byte mask and a value is non-zero, the corresponding value on the attention
            layer will be ignored
        need_weights: output attn_output_weights.
        static_k, static_v: static key and value used for attention operators.
    Shapes for inputs:
        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
            the embedding dimension. :math:`(N, L, E)` if ``batch_first`` is ``True``.
        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
            the embedding dimension. :math:`(N, S, E)` if ``batch_first`` is ``True``.
        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
            the embedding dimension. :math:`(N, S, E)` if ``batch_first`` is ``True``.
        - attn_bias: if a 2D mask: :math:`(L, S)` where L is the target sequence length, S is the
            source sequence length.
            If a 3D mask: :math:`(N\cdot\text{num\_heads}, L, S)` where N is the batch size, L is the target sequence
            length, S is the source sequence length. ``attn_bias`` allows to pass pos embed directly into attention
            If a ByteTensor is provided, the non-zero positions are not allowed to attend while the zero positions will
            be unchanged. If a BoolTensor is provided, positions with ``True`` is not allowed to attend while ``False`` 
            values will be unchanged. If a FloatTensor is provided, it will be added to the attention weight.
        - attn_mask: if a 2D mask: :math:`(L, S)` where L is the target sequence length, S is the
            source sequence length.
            If a 3D mask: :math:`(N\cdot\text{num\_heads}, L, S)` where N is the batch size, L is the target sequence
            length, S is the source sequence length. ``attn_mask`` ensure that position i is allowed to attend
            the unmasked positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
            while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
            is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
            is provided, it will be added to the attention weight.
        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
            If a ByteTensor is provided, the non-zero positions will be ignored while the position
            with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
            value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
            N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
            N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
    Outputs:
        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
            E is the embedding dimension. :math:`(N, L, E)` if ``batch_first`` is ``True``.
        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
            L is the target sequence length, S is the source sequence length.
        """
        if self.batch_first:
            query, key, value = [x.transpose(1, 0) for x in (query, key, value)]

        # set up shape vars
        tgt_len, bsz, embed_dim = query.shape
        src_len, _, _ = key.shape
        if not key.shape[:2] == value.shape[:2]:
            raise ValueError(f"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}")

        q, k, v = self.in_projection(query, key, value)

        # prep attention mask
        if attn_mask is not None:
            if attn_mask.dtype == torch.uint8:
                warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
                attn_mask = attn_mask.to(torch.bool)
            else:
                assert attn_mask.is_floating_point() or attn_mask.dtype == torch.bool, \
                    f"Only float, byte, and bool types are supported for attn_mask, not {attn_mask.dtype}"
            # ensure attn_mask's dim is 3
            if attn_mask.dim() == 2:
                correct_2d_size = (tgt_len, src_len)
                if attn_mask.shape != correct_2d_size:
                    raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
                attn_mask = attn_mask.unsqueeze(0)
            elif attn_mask.dim() == 3:
                correct_3d_size = (bsz * self.num_heads, tgt_len, src_len)
                if attn_mask.shape != correct_3d_size:
                    raise RuntimeError(f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
            else:
                raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")

        # prep key padding mask
        if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
            warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
            key_padding_mask = key_padding_mask.to(torch.bool)

        # add bias along batch dimension (currently second)
        if self.bias_k is not None and self.bias_v is not None:
            assert static_k is None, "bias cannot be added to static key."
            assert static_v is None, "bias cannot be added to static value."
            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
            if attn_mask is not None:
                attn_mask = F.pad(attn_mask, (0, 1))
            if key_padding_mask is not None:
                key_padding_mask = F.pad(key_padding_mask, (0, 1))
        else:
            assert self.bias_k is None
            assert self.bias_v is None

        #
        # reshape q, k, v for multihead attention and make em batch first
        #
        q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        if static_k is None:
            k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        else:
            # TODO finish disentangling control flow so we don't do in-projections when statics are passed
            assert static_k.size(0) == bsz * self.num_heads, \
                f"expecting static_k.size(0) of {bsz * self.num_heads}, but got {static_k.size(0)}"
            assert static_k.size(2) == self.head_dim, \
                f"expecting static_k.size(2) of {self.head_dim}, but got {static_k.size(2)}"
            k = static_k
        if static_v is None:
            v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
        else:
            # TODO finish disentangling control flow so we don't do in-projections when statics are passed
            assert static_v.size(0) == bsz * self.num_heads, \
                f"expecting static_v.size(0) of {bsz * self.num_heads}, but got {static_v.size(0)}"
            assert static_v.size(2) == self.head_dim, \
                f"expecting static_v.size(2) of {self.head_dim}, but got {static_v.size(2)}"
            v = static_v

        # add zero attention along batch dimension (now first)
        if self.add_zero_attn:
            zero_attn_shape = (bsz * self.num_heads, 1, self.head_dim)
            k = torch.cat([k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1)
            v = torch.cat([v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1)
            if attn_mask is not None:
                attn_mask = F.pad(attn_mask, (0, 1))
            if key_padding_mask is not None:
                key_padding_mask = F.pad(key_padding_mask, (0, 1))

        # update source sequence length after adjustments
        src_len = k.size(1)

        # merge key padding and attention masks
        if key_padding_mask is not None:
            assert key_padding_mask.shape == (bsz, src_len), \
                f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
            key_padding_mask = key_padding_mask.view(bsz, 1, 1, src_len).   \
                expand(-1, self.num_heads, -1, -1).reshape(bsz * self.num_heads, 1, src_len)
            if attn_mask is None:
                attn_mask = key_padding_mask
            elif attn_mask.dtype == torch.bool:
                attn_mask = attn_mask.logical_or(key_padding_mask)
            else:
                attn_mask = attn_mask.masked_fill(key_padding_mask, float("-inf"))

        # convert mask to float
        if attn_mask is not None and attn_mask.dtype == torch.bool:
            new_attn_mask = torch.zeros_like(attn_mask, dtype=torch.float)
            new_attn_mask.masked_fill_(attn_mask, float("-inf"))
            attn_mask = new_attn_mask

        #
        # (deep breath) calculate attention and out projection
        #
        attn_output, attn_output_weights = self.attention(q, k, v, attn_bias, attn_mask)
        attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
        attn_output = self.out_projection(attn_output)

        attn_output_weights = attn_output_weights.view(bsz, self.num_heads, tgt_len, src_len) if need_weights else None

        if self.batch_first:
            return attn_output.transpose(1, 0), attn_output_weights
        else:
            return attn_output, attn_output_weights

    def in_projection(self, q: Tensor, k: Tensor, v: Tensor) -> List[Tensor]:
        r"""
        Performs the in-projection step of the attention operation, using packed weights.
        Output is a triple containing projection tensors for query, key and value.
        Args:
            q, k, v: query, key and value tensors to be projected. For self-attention,
                these are typically the same tensor; for encoder-decoder attention,
                k and v are typically the same tensor. (We take advantage of these
                identities for performance if they are present.) Regardless, q, k and v
                must share a common embedding dimension; otherwise their shapes may vary.
        Shape:
            Inputs:
            - q: :math:`(..., E)` where E is the embedding dimension
            - k: :math:`(..., E)` where E is the embedding dimension
            - v: :math:`(..., E)` where E is the embedding dimension
            Output:
            - in output list :math:`[q', k', v']`, each output tensor will have the
                same shape as the corresponding input tensor.
        """
        if k is v:
            # self-attention
            if q is k:
                return self.in_proj(q).split((self.embed_dim, self.k_dim, self.v_dim), dim=-1)
            # encoder-decoder attention
            else:
                w_q, w_kv = self.in_proj.weight.split([self.embed_dim, self.k_dim + self.v_dim])
                b_q, b_kv = None if self.in_proj.bias is None else self.in_proj.bias.split([self.embed_dim, self.k_dim + self.v_dim])
                return (F.linear(q, w_q, b_q),) + F.linear(k, w_kv, b_kv).split((self.k_dim, self.v_dim), dim=-1)
        else:
            w_q, w_k, w_v = self.in_proj.weight.split([self.embed_dim, self.k_dim, self.v_dim])
            b_q, b_k, b_v = None if self.in_proj.bias is None else self.in_proj.bias.split([self.embed_dim, self.k_dim, self.v_dim])
            return F.linear(q, w_q, b_q), F.linear(k, w_k, b_k), F.linear(v, w_v, b_v)

    def attention(self, q: Tensor, k: Tensor, v: Tensor, attn_bias: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None,) -> Tuple[Tensor, Tensor]:
        r"""
        Computes scaled dot product attention on query, key and value tensors, using
        an optional attention mask if passed, and applying dropout if a probability
        greater than 0.0 is specified.
        Returns a tensor pair containing attended values and attention weights.
        Args:
            q, k, v: query, key and value tensors. See Shape section for shape details.
            attn_mask: optional tensor containing mask values to be added to calculated
                attention. May be 2D or 3D; see Shape section for details.
            attn_bias: optional tensor containing bias values to be added to calculated
                attention. Used for relative positional embedding. May be 2D or 3D; see
                Shape section for details.
        Shape:
            - q: :math:`(B, Nt, E)` where B is batch size, Nt is the target sequence length,
                and E is embedding dimension.
            - key: :math:`(B, Ns, E)` where B is batch size, Ns is the source sequence length,
                and E is embedding dimension.
            - value: :math:`(B, Ns, E)` where B is batch size, Ns is the source sequence length,
                and E is embedding dimension.
            - attn_bias: either a 3D tensor of shape :math:`(B, Nt, Ns)` or a 2D tensor of
                shape :math:`(Nt, Ns)`.
            - attn_mask: either a 3D tensor of shape :math:`(B, Nt, Ns)` or a 2D tensor of
                shape :math:`(Nt, Ns)`.
            - Output: attention values have shape :math:`(B, Nt, E)`; attention weights
                have shape :math:`(B, Nt, Ns)`
        """
        q *= self.scaling
        # (B, Nt, E) x (B, E, Ns) -> (B, Nt, Ns)
        attn = torch.bmm(q, k.transpose(-2, -1))
        if attn_bias is not None:
            attn += attn_bias
        if attn_mask is not None:
            attn += attn_mask
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)
        # (B, Nt, Ns) x (B, Ns, E) -> (B, Nt, E)
        output = torch.bmm(attn, v)
        return output, attn

    def out_projection(self, attn_output: Tensor) -> Tensor:
        return self.out_proj(attn_output)


class FeedForwardNetwork(nn.Module):

    def __init__(self, token_size: int, feedforward_size: int,
                 activation: Optional[str] = 'GELU',
                 feedforward_dropout: Optional[float] = 0.,
                 **kwargs: Dict[str, Any]) -> None:
        super(FeedForwardNetwork, self).__init__()
        self.fc1 = nn.Linear(token_size, feedforward_size)
        self.fc2 = nn.Linear(feedforward_size, token_size)
        self.activation = getattr(nn, activation)()
        self.dropout = nn.Dropout(feedforward_dropout)

    def forward(self, x):
        out = self.fc1(x)
        out = self.activation(out)
        out = self.dropout(out)
        out = self.fc2(out)

        return out


class TransformerEncoderLayer(nn.Module):
    r"""TransformerEncoderLayer is made up of self-attn and feedforward network.
    This standard encoder layer is based on the paper "Attention Is All You Need".
    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez,
    Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in
    Neural Information Processing Systems, pages 6000-6010. Users may modify or implement
    in a different way during application.
    Args:
        token_size: the number of expected features in the input (required).
        num_heads: the number of heads in the multiheadattention models (required).
        feedforward_size: the dimension of the feedforward network model (default=2048).
        dropout: the dropout value (default=0.1).
        activation: the activation function of intermediate layer, relu or gelu (default=relu).
        layer_norm_eps: the eps value in layer normalization components (default=1e-5).
        batch_first: If ``True``, then the input and output tensors are provided
            as (batch, seq, feature). Default: ``False``.
        norm_first: if ``True``, layer norm is done prior to attention and feedforward
            operations, respectivaly. Otherwise it's done after. Default: ``False`` (after).
    Examples::
        >>> encoder_layer = nn.TransformerEncoderLayer(token_size=512, num_heads=8)
        >>> src = torch.rand(10, 32, 512)
        >>> out = encoder_layer(src)
    Alternatively, when ``batch_first`` is ``True``:
        >>> encoder_layer = nn.TransformerEncoderLayer(token_size=512, num_heads=8, batch_first=True)
        >>> src = torch.rand(32, 10, 512)
        >>> out = encoder_layer(src)
    """
    __constants__ = ['batch_first', 'norm_first']

    def __init__(self, token_size: int, num_heads: int,
                 feedforward_size: Optional[int] = 2048,
                 dropout: Optional[float] = 0.1,
                 attention_dropout: Optional[float] = 0.1,
                 feedforward_dropout: Optional[float] = 0.1,
                 activation: Optional[str] = 'GELU',
                 layer_norm_eps: Optional[float] = 1e-5,
                 scale_factor: Optional[float] = 1.,
                 bias: Optional[bool] = True,
                 add_bias_kv: Optional[bool] = False,
                 add_zero_attn: Optional[bool] = False,
                 batch_first: Optional[bool] = True,
                 norm_first: Optional[bool] = False,
                 Attention: Optional[Type[nn.Module]] = MultiheadAttention,
                 FeedForwardNetwork: Optional[Type[nn.Module]] = FeedForwardNetwork,
                 **kwargs: Dict[str, Any]) -> None:
        super(TransformerEncoderLayer, self).__init__()
        self.norm_first = norm_first
        self.attention = Attention(
            token_size, num_heads, attention_dropout=attention_dropout,
            scale_factor=scale_factor, bias=bias, add_bias_kv=add_bias_kv,
            add_zero_attn=add_zero_attn, batch_first=batch_first, **kwargs)
        self.norm1 = nn.LayerNorm(token_size, eps=layer_norm_eps)
        self.feedforward = FeedForwardNetwork(
            token_size, feedforward_size, activation, feedforward_dropout, **kwargs)
        self.norm2 = nn.LayerNorm(token_size, eps=layer_norm_eps)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src: Tensor,
                attn_bias: Optional[Tensor] = None,
                attn_mask: Optional[Tensor] = None,
                key_padding_mask: Optional[Tensor] = None,
                need_weights: Optional[bool] = True) -> Tensor:
        r"""Pass the input through the encoder layer.
        Args:
            src: the sequence to the encoder layer (required).
            attn_mask: the mask for the src sequence (optional).
            key_padding_mask: the mask for the src keys per batch (optional).
        Shape:
            see the docs in Transformer class.
        """

        if self.norm_first:
            src = self.norm1(src)

        attn, weights = self.attention(
            src, src, src, attn_bias=attn_bias, attn_mask=attn_mask,
            key_padding_mask=key_padding_mask, need_weights=need_weights)
        attn = src + self.dropout(attn)
        attn = self.norm1(attn) if not self.norm_first else self.norm2(attn)

        ffn = self.feedforward(attn)
        ffn = attn + self.dropout(ffn)
        if not self.norm_first:
            ffn = self.norm2(ffn)

        return ffn, weights


class TransformerEncoder(nn.Module):
    r"""TransformerEncoder is a stack of N encoder layers
    Args:
        layer: the sub-encoder-layer in the encoder (required).
        num_layers: the number of sub-encoder-layers in the encoder (required).
        drop_layer: the drop layer rate (default=0.0).
    Examples::
        >>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
        >>> src = torch.rand(10, 32, 512)
        >>> out = transformer_encoder(src)
    """
    __constants__ = ['norm']

    def __init__(self, num_layers: int,
                 drop_layer: Optional[float] = 0.,
                 norm: Optional[nn.Module] = None,
                 Layer: Optional[Type[nn.Module]] = TransformerEncoderLayer,
                 Attention: Optional[Type[nn.Module]] = MultiheadAttention,
                 FeedForwardNetwork: Optional[Type[nn.Module]] = FeedForwardNetwork,
                 **kwargs: Dict[str, Any]) -> None:
        super(TransformerEncoder, self).__init__()
        self.num_layers = num_layers
        self.layers = nn.ModuleList([])
        self.layers.extend([Layer(Attention=Attention, FeedForwardNetwork=FeedForwardNetwork, **kwargs) for _ in range(self.num_layers)])
        self.norm = norm

    def forward(self, src: Tensor, attn_bias: Optional[Tensor] = None,
                attn_mask: Optional[Tensor] = None,
                key_padding_mask: Optional[Tensor] = None,
                need_weights: Optional[bool] = True) -> Tensor:
        r"""Pass the input through the encoder layers in turn.
        Args:
            src: the sequence to the encoder (required).
            attn_mask: the mask for the src sequence (optional).
            key_padding_mask: the mask for the src keys per batch (optional).
        Shape:
            see the docs in Transformer class.
        """
        output = src
        attn_weights = [] if need_weights else None

        for layer in self.layers:
            output, weights = layer(output, attn_bias=attn_bias, attn_mask=attn_mask, key_padding_mask=key_padding_mask, need_weights=need_weights)
            if need_weights:
                attn_weights.append(weights)

        if need_weights:
            attn_weights = torch.stack(attn_weights).cpu().detach()

        return output, attn_weights
```


cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @bhosmer @cpuhrsch