Convert model from flax to TF
Author: arunraja-hub
Pull request: no

I am trying to convert my flax MT5 model to TensorFlow. I devised the following script using https://github.com/huggingface/transformers/issues/12545

```
from transformers import MT5Model, MT5TokenizerFast, TFMT5Model, MT5Config, FlaxT5ForConditionalGeneration

import numpy as np
import jax
import jax.numpy as jnp

pretrained = "../dumped/code-mt5-large-batch-mix/" # earlier missed the fact that there is no ckpt in this dir
tmp_path = "../dumped/code-mt5-large-batch-mix-tensorflow"

config = MT5Config.from_pretrained(pretrained, from_flax=True)
model = FlaxT5ForConditionalGeneration.from_pretrained(pretrained, config=config)
tokenizer = MT5TokenizerFast(pretrained, use_fast=True, extra_ids=160)

def to_f32(t):
    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)

model.params = to_f32(model.params)
model.save_pretrained(tmp_path)

model_tf = TFMT5Model.from_pretrained(tmp_path)
model_tf.save_pretrained(tmp_path)

```
However, the conversion gets aborted with this output: https://paste.ubuntu.com/p/Ynw9Tn8NC9/

According to the output, the conversion seems to require a specific file instead of the entire model directory `../dumped/code-mt5-large-batch-mix/`  (`  what():  basic_filebuf::underflow error reading the file: Is a directory
`). We are not sure if this is the case and if so what is the specific file required.

The contents of  `../dumped/code-mt5-large-batch-mix/` are:
![image](https://user-images.githubusercontent.com/43485111/125642556-933161e1-a8e3-4611-8097-9ba9f8c0e022.png)

Some help with this model conversion is much appreciated. Thanks!

Comments:
At the moment we only have Flax <=> PT and TF <=> PT conversion. So you should do the following:

```python
from transformers import T5ForConditionalGeneration, MT5TokenizerFast, TFT5ForConditionalGeneration, MT5Config, FlaxT5ForConditionalGeneration

import numpy as np
import jax
import jax.numpy as jnp

pretrained = "../dumped/code-mt5-large-batch-mix/" # earlier missed the fact that there is no ckpt in this dir
tmp_path = "../dumped/code-mt5-large-batch-mix-tensorflow"

config = MT5Config.from_pretrained(pretrained, from_flax=True)
model = T5ForConditionalGeneration.from_pretrained(pretrained, config=config)
tokenizer = MT5TokenizerFast(pretrained, use_fast=True, extra_ids=160)

def to_f32(t):
    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)

model.params = to_f32(model.params)
model.save_pretrained(tmp_path)

model_pt = T5ForConditionalGeneration.from_pretrained(tmp_path, from_flax=True)
model_pt.save_pretrained(tmp_path)

model_tf = TFT5ForConditionalGeneration.from_pretrained(tmp_path, from_pt=True)
model_tf.save_pretrained(tmp_path)
```

Thank you!

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.