Unable to import name 'pad_shard_unpad' from 'flax.jax_utils' (clm language modelling flax example)
Author: SupreethRao99
Pull request: no

### System Info

- `transformers` version: 4.26.0.dev0
- Platform: Linux-5.4.88+-x86_64-with-glibc2.2.5
- Python version: 3.8.15
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 1.11.0+cu102 (False)
- Tensorflow version (GPU?): 2.10.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.4.2 (tpu)
- Jax version: 0.3.10
- JaxLib version: 0.3.10
- Flax version: 0.4.2
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: Yes (TPUv3-8 1VM Kaggle)

### Who can help?

@sanchit-gandhi @ArthurZucker @younesbelkada @sgugger 

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I've been trying to run `transformers/examples/flax/language-modeling/run_clm_flax.py` on Kaggle's new TPUv3-8 1VM type of instance. it's a TPU instance with the TPU devices directly attached.

When I ran the example code in `transformers/examples/flax/language-modeling` for causal language modelling with the following code

```
!python /kaggle/working/transformers/examples/flax/language-modeling/run_clm_flax.py \
    --output_dir="<models_direcotry>" \
    --model_type="gpt2" \
    --config_name="<custom_GPT_Config>" \
    --tokenizer_name="<custom_tokenizer>" \
    --dataset_name="<path_to_dataset_on_hf_hub>" \
    --do_train --do_eval \
    --block_size="128" \
    --per_device_train_batch_size="64" \
    --per_device_eval_batch_size="64" \
    --learning_rate="5e-3" --warmup_steps="1000" \
    --adam_beta1="0.9" --adam_beta2="0.98" --weight_decay="0.01" \
    --overwrite_output_dir \
    --num_train_epochs="20" \
    --logging_steps="500" \
    --save_steps="2500" \
    --eval_steps="2500"
```
I'm getting the following error 
```
WARNING: Logging before InitGoogle() is written to STDERR
I0000 00:00:1671849395.203410    2540 tpu_initializer_helper.cc:116] libtpu.so is already in use by process with pid 12. Not attempting to load libtpu.so in this process.
2022-12-24 02:36:35.892840: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2022-12-24 02:36:35.936831: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2022-12-24 02:36:36.693369: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2022-12-24 02:36:36.693463: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2022-12-24 02:36:36.693483: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
Traceback (most recent call last):
  File "/kaggle/working/transformers/examples/flax/language-modeling/run_clm_flax.py", line 46, in <module>
    from flax.jax_utils import pad_shard_unpad, unreplicate
ImportError: cannot import name 'pad_shard_unpad' from 'flax.jax_utils' (/usr/local/lib/python3.8/site-packages/flax/jax_utils.py)
```

I can't upgrade jax, jaxlib, flax as it messes with the TPU's connected causing them to become unavailable.

### Expected behavior

GPT2 training should begin on 8 TPUv3 devices

Comments:
@sanchit-gandhi @ArthurZucker, could you help us with this ^

Hi @SupreethRao99 

Indeed the function `pad_shard_unpad` cannot be imported from flax.jax_utils using flax==0.4.2
Can you try with a latest version of `flax`? For example `from flax.jax_utils import pad_shard_unpad` works fine under `flax==0.5.3` | `pip install --upgrade flax` or `pip install flax==0.5.3`

Hi @younesbelkada , Upgrading flax to the latest version caused some issues jax and jaxlib but using `flax==0.5.3` along with the latest kaggle runtime (30-12-2022)  fixed the issue. Thanks a lot !