Repository: pytorch
Issue: Need TransformerEncoder to output attention map
Labels: module: nn,triaged,topic: new features

### 🚀 The feature, motivation and pitch

I need to observe the strength of attention for each element in the sequence with all the elements in the same sequence. This is called the attention map in the language of the transformer and is one of the important metrics for visualization. I want the output to be of shape num_layers x [batch_size, sequence_length, sequence_length]. 

The MultiheadAttention layer can return the attention map when `need_weights` is provided in the forward call, however, it is always set to `false` from the forward call of layer `TransformerEncoder`. It would be better if we could specify `need_weights` in `TransformerEncoder`'s forward call.

## Proposed working changes:
File:
[torch/nn/modules/transformer.py](https://github.com/pytorch/pytorch/blob/148d49260aa29ba6d4c8354a019f64a27503d5b8/torch/nn/modules/transformer.py)
### Before:
```
...
class TransformerEncoder(Module):
...
    def forward(
            self,
            src: Tensor,
            mask: Optional[Tensor] = None,
            src_key_padding_mask: Optional[Tensor] = None,
            is_causal: Optional[bool] = None) -> Tensor:
...
        for mod in self.layers:
            output = mod(output, src_mask=mask, is_causal=is_causal, src_key_padding_mask=src_key_padding_mask_for_layers)
...
        return output
...
class TransformerEncoderLayer(Module):
...
        x = src
        if self.norm_first:
            x  = x +  self._sa_block(self.norm1(x), src_mask, src_key_padding_mask)
            x = x + self._ff_block(self.norm2(x))
        else:
            x = self.norm1(x + self._sa_block(x, src_mask, src_key_padding_mask))
            x = self.norm2(x + self._ff_block(x))

        return x

    # self-attention block
    def _sa_block(self, x: Tensor,
                  attn_mask: Optional[Tensor], key_padding_mask: Optional[Tensor]) -> Tensor:
        x = self.self_attn(x, x, x,
                           attn_mask=attn_mask,
                           key_padding_mask=key_padding_mask,
                           need_weights=False)[0]
        return self.dropout1(x)
...
```

### After:
```
...
class TransformerEncoder(Module):
...
    def forward(
            self,
            src: Tensor,
            mask: Optional[Tensor] = None,
            src_key_padding_mask: Optional[Tensor] = None,
            is_causal: Optional[bool] = None,
            need_weights: Optional[bool] = False) -> Tensor:
...
        attn_maps = []
        for mod in self.layers:
            output, attn_map = mod(output, src_mask=mask, is_causal=is_causal, src_key_padding_mask=src_key_padding_mask_for_layers, need_weights=need_weights)
            attn_maps.append(attn_map)
...
        return output, attn_maps
...
class TransformerEncoderLayer(Module):
...
        x = src
        if self.norm_first:
            _x, attn_map = self._sa_block(self.norm1(x), src_mask, src_key_padding_mask, need_weights)
            x  = x + _x
            x = x + self._ff_block(self.norm2(x))
        else:
            _x, attn_map = self._sa_block(x, src_mask, src_key_padding_mask, need_weights)
            x = self.norm1(x + _x)
            x = self.norm2(x + self._ff_block(x))

        return x, attn_map

    # self-attention block
    def _sa_block(self, x: Tensor,
                  attn_mask: Optional[Tensor], key_padding_mask: Optional[Tensor], need_weights: Optional[bool] = False) -> Tensor:
        x, attn_map = self.self_attn(x, x, x,
                           attn_mask=attn_mask,
                           key_padding_mask=key_padding_mask,
                           need_weights=need_weights)
        return self.dropout1(x), attn_map
...
```

### Additional context

Pytorch version: 2.0.0

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @bhosmer @cpuhrsch @erichan1 @drisspg