Automatic Mixed Precision Support for (some) Flax Transformers
# 🚀 Feature request

<!-- A clear and concise description of the feature proposal.
     Please provide a link to the paper and code in case they exist. -->

[Mixed precision](https://arxiv.org/abs/1710.03740) can reduce memory usage and improve computation efficiency. Thanks to bfloat16, it is very straightforward on TPU. However, on GPU it requires type casting for certain operations. It will be great if automatic type casting for mixed precision is implemented in (some popular) Flax models so that they can run efficiently also on GPUs.

## Motivation

Ideally, we would want JAX/Flax Transformers to run easily and efficiently on both GPU and TPU. Having mixed precision will help speed up both training and inference by big margins. However, on GPU this requires both loss/gradient scaling and auto casting implemented. Currently auto casting is not supported by JAX or Flax (although gradient scaling is. )

JAX team said AMP should be external to JAX https://github.com/google/jax/issues/5257.
Flax has a milestone, which is of low priority for the moment https://github.com/google/flax/milestone/6.

My intuition, which may be wrong, is that since we are dealing with transformers, there should be a relatively smaller set operations, such as softmax, that needs special treatments compared to the [full set of ops in pytorch](https://pytorch.org/docs/stable/amp.html#autocast-op-reference).
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.