Attention docstring missing head dimension for arguments mask and nonbatched_bias

Linked to the following [PR](https://github.com/google-deepmind/alphafold/pull/893) and [fork](https://github.com/louis-rf/alphafold).

The Attention module is only called in four places, with the following shapes:
```
- TriangleAttention: [q_data: (B, Q, ..); m_data: (B, K, ..); mask: (B, H=1, Q=1, K); nonbatched_bias: (H, Q, K)]
- TemplateEmbedding: [q_data: (B, Q, ..); m_data: (B, K, ..); mask: (B, H=1, Q=1, K)]
- MSARowAttentionWithPairBias: [q_data: (B, Q, ..); m_data: (B, K, ..); mask: (B, H=1, Q=1, K); nonbatched_bias: (H, Q, K)]
- MSAColumnAttention: [q_data: (B, Q, ..); m_data: (B, K, ..); mask: (B, H=1, Q=1, K)]
```
The mask always has new axes for dimensions 2 and 3. In all four cases it is wrapped in alphafold.model.mapping.inference_subbatch, but this doesn't affect the dimensions, only the sizes.

The docstring gives incorrect shapes for the arguments: mask, nonbatched_bias, based on the usage it should be:
```
mask: A mask for the attention, shape [batch_size, N_heads, N_queries, N_keys].
nonbatched_bias: Shared bias, shape [N_heads, N_queries, N_keys].
```
instead of:
```
mask: A mask for the attention, shape [batch_size, N_queries, N_keys].
nonbatched_bias: Shared bias, shape [N_queries, N_keys].
```
This is clear when looking at where mask is used:
```
...

logits = jnp.einsum('bqhc,bkhc->bhqk', q, k)
if nonbatched_bias is not None:
logits += jnp.expand_dims(nonbatched_bias, axis=0)
logits = jnp.where(mask, logits, _SOFTMAX_MASK)
...
```
as the output of the einsum has shape: `bhqk`.

I believe some implementations of attention wont have a head dimension in the mask, since it is not used in AlphaFold maybe it would be worth removing it in the mask when attention is called (and including an expand_dims for this head dimension within the attention module). But only changing the docstring is easier, and it is still a valid implementation of Attention so I think that is the way to go.