[Flax] Bump to v0.4.1

[sanchit-gandhi]

What does this PR do?

The flatten_dict operator with the kwarg argument sep was added to modeling_flax_utils in #17760:

transformers/src/transformers/modeling_flax_utils.py

Line 127 in f25457b

weights = flatten_dict(params, sep="/")

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).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• 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

[ArthurZucker]

Nice finding! I am not competent regarding which versions we want to always support, but LGTM

[HuggingFaceDocBuilderDev]

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

[patil-suraj]

Good catch! I'm fine with always using the latest flax version

[sanchit-gandhi]

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?

[patrickvonplaten]

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

[patrickvonplaten]

Ok to bump out of necessity (nice find @sanchit-gandhi) - not because it's best to force the user to always use the newest version. E.g. one might very well have an old research project that only works with Flax 0.3.5, but would like to use the newest transformers. Now, the person could not combine the two -> let's try to stay compatible with older versions in the future :-)

[sanchit-gandhi]

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!

[patrickvonplaten]

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