[Flax] Bump to v0.4.1
# What does this PR do?

The `flatten_dict` operator with the kwarg argument `sep` was added to `modeling_flax_utils` in https://github.com/huggingface/transformers/pull/17760:
https://github.com/huggingface/transformers/blob/f25457b273348733bfeb19a51ab0d21bd30a08b8/src/transformers/modeling_flax_utils.py#L127

This kwarg was only added to Flax in v0.4.1: https://github.com/google/flax/releases/tag/v0.4.1

This PR bumps the required Flax version in Transformers from v0.3.5 to v0.4.1.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

cc @ArthurZucker 

Nice finding! I am not competent regarding which versions we want to always support, but LGTM<|||||>_The documentation is not available anymore as the PR was closed or merged._<|||||>Thanks @patil-suraj! The latest version would be v0.5.2 (see https://github.com/google/flax/releases). Is the best way to test this by setting `flax>=0.5.2` and observe if the tests past?<|||||>IMO we should also try to stay compatible with older Flax versions, so let's not go as high as we can but only as high as we have to! PR looks good to me - but let's try to not always use the most recent Flax features in order to stay compatible with older versions as well <|||||>Thanks for the details @patrickvonplaten. I'm in agreement that we should avoid potentially breaking use cases purely for the sake of being on the latest version, but try and integrate the latest features where applicable.

Interestingly, I actually went away and had a play with three different versions of Flax on my personal research project https://github.com/sanchit-gandhi/seq2seq-speech:

1. **v0.3.5:** issues regarding the `sep` arg of `flatten_dict` in `modeling_flax_utils` (described above)
2. **v0.4.2:** no apparent issues
3. **v0.5.2:** (latest) had issues with Flax's `scan` not working depending on JAX version

Seems like v0.4.2 sits in the sweet spot for newer Flax version whilst providing backwards compatibility!<|||||>Yes, but it's often worth to also not directly implement all the new features of Flax since:
- a) they might not work very well because they are new
- b) it breaks backwards comp