pytest jax,jaxlib,flax versions incompatibility
### System Info

When running pytest (```pytest --collect-only -q``` is enough), it fails due to
```
AttributeError: module 'jax.tree_util' has no attribute 'register_pytree_with_keys_class'
```

From setup.py deps:
```
    "jax>=0.2.8,!=0.3.2,<=0.3.6",
    "jaxlib>=0.1.65,<=0.3.6",**
```

## Fix:
pip install jax==0.3.25 jaxlib==0.3.25 flax==0.6.2

### Who can help?

@sanchit-gandhi

### Information

- [X] The official example scripts
- [ ] 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

pytest --collect-only ./tests/models/<some_model>

### Expected behavior

"tests collected" without any error
Could you try:
```
pip install -e ".[flax]"
```
To get the correct version of all three packages?<|||||>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/main/CONTRIBUTING.md) are likely to be ignored.<|||||>More detailed instructions for installing Flax & Transformers can be found here: https://github.com/huggingface/transformers/tree/main/examples/research_projects/jax-projects#how-to-install-relevant-libraries

Leaving this one closed for now. Feel free to open if the above guide doesn't solve your issue and I can take another look