Pin flax & optax version
# What does this PR do?

Failing on [main](https://app.circleci.com/pipelines/github/huggingface/transformers/62534/workflows/d270e074-306d-4a8f-9434-fcdd979fae1b/jobs/770753) because of a new release of [optax](https://github.com/deepmind/optax/releases/tag/v0.1.5). Pinning until compatible versions with jax resolved. 



Fixes # (issue)

_The documentation is not available anymore as the PR was closed or merged._