Incorrect docstring / documentation for torch.nn.functional.scaled_dot_product_attention in 2.1
Labels: module: docs, triaged, module: multi-headed-attention

### 📚 The doc issue

I originally noticed this error from following the documentation page at:
'https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html#torch.nn.functional.scaled_dot_product_attention'

where I noticed that the written 'efficient implementation' was not using the mask correctly:

```
# Efficient implementation equivalent to the following:
def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None) -> torch.Tensor:
    # Efficient implementation equivalent to the following:
    L, S = query.size(-2), key.size(-2)
    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
    attn_bias = torch.zeros(L, S, dtype=query.dtype)
    if is_causal:
        assert attn_mask is None
        temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
        attn_bias.to(query.dtype)

    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_mask.masked_fill_(attn_mask.logical_not(), float("-inf"))
        else:
            attn_bias += attn_mask
    attn_weight = query @ key.transpose(-2, -1) * scale_factor
    attn_weight += attn_bias
    attn_weight = torch.softmax(attn_weight, dim=-1)
    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
    return attn_weight @ value
```

The error is that if the attn_mask.dtype == torch.bool, the attn_mask does not impact the attn_weight as it does not change the attn_bias and also is not mentioned in the final 5 lines. This documentation page is derived from the docstring in the code that has this error.

If we look at the current documentation page representing the docstring for the unstable version of the package at: https://pytorch.org/docs/main/generated/torch.nn.functional.scaled_dot_product_attention.html#torch.nn.functional.scaled_dot_product_attention, we can see:

```
 # Efficient implementation equivalent to the following:
def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None) -> torch.Tensor:
    # Efficient implementation equivalent to the following:
    L, S = query.size(-2), key.size(-2)
    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
    attn_bias = torch.zeros(L, S, dtype=query.dtype)
    if is_causal:
        assert attn_mask is None
        temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
        attn_bias.to(query.dtype)

    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
        else:
            attn_bias += attn_mask
    attn_weight = query @ key.transpose(-2, -1) * scale_factor
    attn_weight += attn_bias
    attn_weight = torch.softmax(attn_weight, dim=-1)
    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
    return attn_weight @ value
```

where notably they masked_fill attn_bias instead of attn_mask.

N.B: I haven't really spent the time checking whether the code is implemented as in the docstring: READ: Incorrectly. 

### Suggest a potential alternative/fix

A fix is clear. Change the docstring that is used in V2.1 so as to have attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf")) as in the unstable version
rather than attn_mask.masked_fill_(attn_mask.logical_not(), float("-inf")).

It may be worth considering whether this change can be explicitly made to the online documentation before doing a v2.1.1 where this error is amended.



cc @svekars @carljparker