Update flax requirement from <0.7,>=0.6.5 to >=0.6.5,<0.8

Created by: dependabot[bot]

Updates the requirements on flax to permit the latest version.

Release notes

Sourced from flax's releases.

Version 0.7.0

What's Changed

• delete long-unused dotgetter utility by @​copybara-service in google/flax#3156
• Add sections to flax.linen toctree by @​cgarciae in google/flax#3073
• update doctest requirements to use recentish jax minver by @​levskaya in google/flax#3160
• Update gym to gymnasium by @​cgarciae in google/flax#3133
• Add lifted transforms section to Haiku migration guide by @​cgarciae in google/flax#3158
• relaxed tolerance by @​chiamp in google/flax#3162
• fixed broken links by @​chiamp in google/flax#3161
• added dict migration guide by @​chiamp in google/flax#3109
• Update requirements, restructure files and fix formatting for VAE example by @​canyon289 in google/flax#3046
• Update python version support by @​cgarciae in google/flax#3168
• Improve Haiku migration guide by @​cgarciae in google/flax#3169
• Set default types in Flax for Orbax restoration and add restore_with_serialized_types in preparation for an upcoming change. by @​copybara-service in google/flax#3165
• fix precommit issues by @​chiamp in google/flax#3170
• fixed incorrect reference link by @​chiamp in google/flax#3167
• added absltest to linen_recurrent_test by @​chiamp in google/flax#3172
• remove cell_size docs from RNN by @​cgarciae in google/flax#3186
• 0.7.0 by @​cgarciae in google/flax#3187

Full Changelog: https://github.com/google/flax/compare/v0.6.11...v0.7.0

Changelog

Sourced from flax's changelog.

0.7.0

• RNNCellBase refactor.

0.6.11

• Set Orbax-as-backend to be the default checkpointing method.
• Fix setup trigger issue under sharing and transforms.
• Add collection to self.scope.reserve(name, col) so that sow works with the same name in different collections.
• Minor improvements for Sequential.
• Improve the error message in MultiHeadDotProductAttention.
• Allow manually specifying the rng key for Dropout.
• RNN refactor.
• fixed missing separator for rng fold in.

0.6.10

• Rudimentary quantization support: some layers can be parametrized with custom dot_general and conv_general_dilated.

0.6.9

• Depend on orbax-checkpoint package instead of orbax.
• Refactored setup scripts to project.toml.
• Added pretty_repr utility fn.
• Fix get_partition_spec on replicated array.
• Updates imagenet.ipynb to use GPU Colab runtime, and fixed config.
• Upgrade checkpointing code to jax.sharding, and with more warnings.

0.6.8

• The automatic checkpoint migration was temporarily rolled back due to legacy compatibility issues.
  • We still recommend you to use the upgrade guide and migrate completely to the Orbax API to ensure stability.
  • Or alternatively, add flax.config.update('flax_use_orbax_checkpointing', True) to your project to avoid being impacted by the automatic migration process.
• Added utility functions to frozen_dict api.
• Migrated Flax away from register_keypaths.
• Fixes kwargs in convert_to_graphs_tuple_fn.
• Fixed examples in a few ways:
  • Bumped the TF version
  • Used latest checkpoint formats
  • Other misc fixes.

0.6.7

• New checkpoints will be saved using Orbax! Please check out upgrade guide and consider migrating completely to the Orbax API.
  • You could flax.config.update('flax_use_orbax_checkpointing', False) to temporarily disable this migration, but note that Flax legacy checkpointing will be removed 3 months from Mar 10, 2023.
• Migrating FrozenDict to regular dict: utility functions now work on both.
• Migrated Flax dataclass and FrozenDict to JAX pytree keypath API.
• Fixed pytype and improved typing for Module
• Fixed up uses of PyTree and PyTreeDef types.

... (truncated)

Commits

• 7e9d708 Merge pull request #3187 from google:0.7.0
• eb1c5ee Merge pull request #3186 from google:fix-rnn-docs
• d405953 fix changelog
• 047e4e8 0.7.0
• aac2f05 remove cell_size docs from RNN
• 0734d00 [linen] Only pmean the mean if it is non-zero, in normalization modules.
• 1e7a8b1 [linen] Minor cleanup to normalization code.
• 1942f3b Add RMSNorm to documentation
• 8efa997 Add propagation of param_dtype to carry initializer.
• b05c673 Merge pull request #3172 from chiamp:add_absl
• Additional commits viewable in compare view

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

---

Dependabot commands and options

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot merge will merge this PR after your CI passes on it
• @dependabot squash and merge will squash and merge this PR after your CI passes on it
• @dependabot cancel merge will cancel a previously requested merge and block automerging
• @dependabot reopen will reopen this PR if it is closed
• @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)