Suboptimal default initialization of q/k/v projections in `nn.MultiHeadDotProductAttention`

Initialization of q/k/v projections are not forward-backward normalized for [`linen.MultiHeadDotProductAttention`](https://flax.readthedocs.io/en/v0.8.0/_modules/flax/linen/attention.html#MultiHeadDotProductAttention). This implementation does not face optimization issues when a `pre-LN` variant of transformer is used; but faces convergence issues in the vanilla `post-LN` variant from "Attention is All You Need".

### System information
- Flax version: `flax==0.8.4`

### Problem you have encountered:
Most papers over the past 4 years still use the vanilla `post-LN` transformer. One such is [facebookresearch/detr](https://github.com/facebookresearch/detr). Inputs to the self-attention block on the first decoder is all `zeros`, as shown below:

```python
# src: https://github.com/facebookresearch/detr/blob/29901c51d7fe8712168b8d0d64351170bc0f83e0/models/transformer.py#L55
...
tgt = torch.zeros_like(query_embed)
memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
pos=pos_embed, query_pos=query_embed)
...
```
The default `linen.MultiHeadDotProductAttention` does not converge at all for the first 10 epochs. The reason is initialization and gradient norm behavior of q/k/v matrices, particularly when inputs to the MHDPA block are all zeros. I have a diff [here](https://github.com/MasterSkepticista/detr/commit/995f335237b72cf17fb1e187b8cb6faf5d51e784) which converges from the very first epoch.

**Please note:** I have isolated the problem to be in the kernel initializer alone, which this proposal is about. Custom implementation of MHDPA does not have any impact on this proposal.

This convergence issue does not exist in PyTorch, because q/k/v are:
1. `xavier_uniform` initialized, and
2. with appropriate `fan_in` and `fan_out` values.

### What you expected to happen:
Expected `linen.MultiHeadDotProductAttention` to converge from the first epoch, as is the case in [facebookresearch/detr](https://github.com/facebookresearch/detr)

### Fix:
1. Switch to `xavier_uniform` initializer for projections (versus `default_kernel_init` which is `lecun_normal`). This is also standard best practice. [t-fixup paper](https://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf)
3. Use correct `fan_in` value for the initializer (for same embedding dimensions of q/k/v, `fan_in` should be `3 * embed_dim`) [ref1](https://github.com/pytorch/pytorch/blob/533c4190f961ca969e7498bfaaad85db3050ff2e/torch/nn/modules/activation.py#L1087), [ref2](https://github.com/pytorch/pytorch/blob/533c4190f961ca969e7498bfaaad85db3050ff2e/torch/nn/modules/activation.py#L1113)

Here is an approximate diff needed for change in Flax:

https://github.com/MasterSkepticista/detr/commit/995f335237b72cf17fb1e187b8cb6faf5d51e784

There could be multiple ways to go about calculating correct `fan_in`, or using a giant `dense` layer as PyTorch does it. This has performance implications.

Let me know if the fix in my code (barring the hardcoded values) is a reasonable approach?

Happy to do a PR.