[FLAX] Add dtype to embedding for bert/bart/opt/t5
## What does this PR do?

This PR is the follow-up of #18462. It adds dtype to `nn.Embed` for more common Flax models, including bert, bart, opt, t5, and their copies.
This dtype is necessary for mixed precision training.

## Who can review?
@patrickvonplaten, @LysandreJik, @sanchit-gandhi 
_The documentation is not available anymore as the PR was closed or merged._<|||||>cc @sanchit-gandhi @younesbelkada @ArthurZucker here<|||||>Can we merge this @ArthurZucker ? How to check whether the slow tests pass?<|||||>Thanks @merrymercy !
You can run slow tests for `t5` for example by running: `RUN_SLOW=1 pytest tests/models/t5/test_modeling_flax_t5.py` - regarding what has been suggested by @sanchit-gandhi , you can just add a new test that initializes a model let's say in `bf16` and tests if the generated sequence is the same than the one that is expected:
```
    @slow
    def test_small_generation_bf16(self):
        model = FlaxT5ForConditionalGeneration.from_pretrained("t5-small", dtype=jnp.bfloat16)
        EXPECTED_OUTPUT = "XXXX"
        self.assertTrue(model.params["shared"]["embedding"].dtype == jnp.bfloat16)
        model.config.max_length = 8
        model.config.num_beams = 1
        model.config.do_sample = False
        tokenizer = T5Tokenizer.from_pretrained("t5-small")

        input_ids = tokenizer("summarize: Hello there", return_tensors="np").input_ids

        sequences = model.generate(input_ids).sequences

        output_str = tokenizer.batch_decode(sequences, skip_special_tokens=True)[0]
        self.assertTrue(output_str == EXPECTED_OUTPUT)

```<|||||>A bf16 test case is added as suggested by @younesbelkada . I checked that slow tests are also passed because most slow tests are in fp32 and this PR does not change the behavior of any fp32 tests.

However, I do notice some slow tests fail on my machine (V100, jax=0.3.25, flax=0.6.2). I think this is not related to my PR, because they fail even with the transformers main branch and transformers v4.24.0. Fixing them is outside the scope of this PR. I can confirm that all tests passed on the main branch can still be passed after my PR.

Since we get two approvals and the test case is added, can we merge it now?<|||||>Thanks again for your contribution!