Missing T5X module 
Author: akku779
Pull request: no

### System Info

I am using the T5X to Pytorch conversion script located in the transformers library to convert my pre-trained T5X model into a Pytorch model, however,  upon running the script I receive the error below.

```
Traceback (most recent call last):
 File "/usr/local/lib/python3.10/dist-packages/transformers/models/t5/convert_t5x_checkpoint_to_pytorch.py", line 36, in <module>
 from t5x import checkpoints
ModuleNotFoundError: No module named 't5x'
```
I have installed the necessary libraries by executing these statements. 
```
!pip install transformers 
!pip install git+https://github.com/google-research/t5x
```

### Who can help?

_No response_

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

The official guidelines for the script is located in transformers/models/t5/convert_t5x_checkpoint_to_pytorch.py
To execute the script, the command below is run. 
```
python3 [path_to_file]/convert_t5x_checkpoint_to_pytorch.py --t5x_checkpoint_path=$HOME/t5_1_1_small --config_file=config.json --pytorch_dump_path=$HOME/t5_1_1_small_pt
```

Where config.json is a config for t5-small (https://huggingface.co/t5-small/blob/main/config.json)



### Expected behavior

The script should convert the checkpoint to a Pytorch model.

Comments:
Hi @akku779, thanks for raising this issue. 

It seems that this is an issue with the installing of the t5x library, rather than one relating to transformers. Running the installation steps I was able to import `t5x` in a python session. 

Given the `!` at the start of the pip commands, were these steps being run in a notebook or ipython environment? In which case, it's necessary to restart to environment in order for the updates to take affect. 

@amyeroberts Thanks for getting back to me. I tried re-running my Colab notebook and I am not receiving the same error. Now, there seems to be some sort of dependency mismatch. 

```
Traceback (most recent call last):
  File "/content/convert_t5x_checkpoint_to_pytorch.py", line 36, in <module>
    from t5x import checkpoints
  File "/usr/local/lib/python3.10/dist-packages/t5x/__init__.py", line 17, in <module>
    import t5x.adafactor
  File "/usr/local/lib/python3.10/dist-packages/t5x/adafactor.py", line 64, in <module>
    from t5x import utils
  File "/usr/local/lib/python3.10/dist-packages/t5x/utils.py", line 43, in <module>
    import orbax.checkpoint
  File "/usr/local/lib/python3.10/dist-packages/orbax/checkpoint/__init__.py", line 20, in <module>
    from orbax.checkpoint import checkpoint_utils
  File "/usr/local/lib/python3.10/dist-packages/orbax/checkpoint/checkpoint_utils.py", line 25, in <module>
    from orbax.checkpoint import type_handlers
  File "/usr/local/lib/python3.10/dist-packages/orbax/checkpoint/type_handlers.py", line 25, in <module>
    from jax.experimental.gda_serialization import serialization
ModuleNotFoundError: No module named 'jax.experimental.gda_serialization'
```

@akku779 From the traceback, this error is coming from the t5x module, and so isn't a transformers issue. I looks like there's a mismatch in your environment between the jax version installed and what the t5x library expects. 

Fixed error by upgrading orbax