In `MultiHeadAttention`, let `num_heads=1` by default
Open
Priority: P2 - no schedule
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
Minor suggestion: In [`nn.MultiHeadAttention`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.MultiHeadAttention), let `num_heads=1` by default, which yields single-head attention. This seems like the most natural default value. I can submit a PR for this.
Contributor guide
Assessment
This issue has not been assessed yet.