Update flax requirement from <0.6,>=0.5 to >=0.5,<0.7

Created by: dependabot[bot]

Updates the requirements on flax to permit the latest version.

Release notes

Sourced from flax's releases.

Version 0.6.0

What's Changed

• Add on_commit_callback to put the responsibility of renaming the directories on the users of the serialization library. This will also fix the GCS atomic rename issue where the users can write a success file when the commit is successful and check the existence of that file before deserialization. by @​copybara-service in google/flax#2328
• RDT Redesign by @​cgarciae in google/flax#2177
• Fix stale URLs to read the docs site. by @​levskaya in google/flax#2338
• Replace all jax.tree_* calls with jax.tree_util.tree_* by @​levskaya in google/flax#2325
• Further fix the singular-leaf checkpointing, and add tests. by @​copybara-service in google/flax#2336
• Forward unroll argument in scan_with_axes by @​sanchit-gandhi in google/flax#2339
• Document repo analytics by @​cgarciae in google/flax#2317
• Improve flax basics by @​cgarciae in google/flax#2291
• Split build into multiple jobs by @​cgarciae in google/flax#2277
• Fix type annotation for step in training.checkpoints by @​Chuxiaof in google/flax#2343
• Add test for writing and restoring empty checkpoints. by @​copybara-service in google/flax#2345
• Allow all processes to checkpoint when not using GDA by @​copybara-service in google/flax#2350
• Make importing tensorstore optional and move related type hints to comments. by @​copybara-service in google/flax#2348
• Use jax.named_scope for name stack rather than named_call. by @​copybara-service in google/flax#2349
• Internal change by @​copybara-service in google/flax#2356
• Fix sphinx CI errors by @​cgarciae in google/flax#2361
• Forward path to rewound Scope by @​jheek in google/flax#2360
• Make link a link on the getting started by @​Davidnet in google/flax#2340
• Fix colab & github links by @​cgarciae in google/flax#2363
• Fix ConvTranspose with circular padding by @​cgarciae in google/flax#2364
• Add some docstrings to the old flax.training.common_utils module. by @​levskaya in google/flax#2373
• Correct state variable name by @​Jeevesh8 in google/flax#2369
• Allow linen's Conv layer to operate on arbitrary-rank inputs. by @​copybara-service in google/flax#2308
• Copies dynamic_scale.py from optim/ to training/. by @​copybara-service in google/flax#2375
• Add option auto_flush to flax.metrics.tensorboard.SummaryWriter by @​copybara-service in google/flax#2376
• updated supported transforms in lifting docs. by @​levskaya in google/flax#2374
• Test docstrings with autodoc on CI by @​cgarciae in google/flax#2372
• skip remat test that fails with autodiff by @​mattjj in google/flax#2389
• Removes flax.optim.dynamic_scale. by @​copybara-service in google/flax#2314
• Plumb spmd_axis_name from vmap_with_axes through to JAX vmap by @​copybara-service in google/flax#2390
• fixed math expressions by @​banda-larga in google/flax#2392

New Contributors

• @​sanchit-gandhi made their first contribution in google/flax#2339
• @​Chuxiaof made their first contribution in google/flax#2343
• @​Davidnet made their first contribution in google/flax#2340
• @​Jeevesh8 made their first contribution in google/flax#2369
• @​banda-larga made their first contribution in google/flax#2392

Full Changelog: https://github.com/google/flax/compare/v0.5.3...v0.6.0

Changelog

Sourced from flax's changelog.

0.6.0

• Removed deprecated optimizers in flax.optim package.
• Moved flax.optim.dynamic_scale to flax.training.dynamic_scale.
• Switched to using jax.named_scope for all profile naming, cut some pointless stack traces out.

0.5.3

New features:

• Added nn.switch as a lifted version of jax.lax.switch.
• Added a method for detecting the use of "init" functions.
• Added checkpointing support for jax.experimental.GlobalDeviceArray, a useful array type for multiprocess/multihost computing.
• Added async option to save_checkpoints() on single-process scenario.
• Improved documentation pages.

Bug fixes:

• Fixed variable aliasing in put_variable
• Fixed missing passthrough of nn.scan unroll arg
• Fixed the MNIST example

0.5.2

• Fixes missing PyYAML dependecy.

0.5.1

New features:

• Added nn.tabulate and Module.tabulate to generate rich representations of the network structure.

0.5.0

• Added flax.jax_utils.ad_shard_unpad() by @​lucasb-eyer
• Implemented default dtype FLIP. This means the default dtype is now inferred from inputs and params rather than being hard-coded to float32. This is especially useful for dealing with complex numbers because the standard Modules will no longer truncate complex numbers to their real component by default. Instead the complex dtype is preserved by default.

Bug fixes:

• Fix support for JAX's experimental_name_stack.

Breaking changes:

• In rare cases the dtype of a layer can change due to default dtype FLIP. See the "Backward compatibility" section of the proposal for more information.

0.4.2

New features:

... (truncated)

Commits

• cda7a4c Removes flax.optim & bumps version to 0.6.0
• 47418b3 Merge pull request #2392 from banda-larga:patch-1
• 99294fb Fix minor typos in struct.py docstring.
• 43c3c29 fixed jax_for_the_impatient line
• 8fe8f5e Plumb spmd_axis_name from vmap_with_axes through to JAX vmap
• d9897c5 Update minimal JAX version to latest (0.3.16).
• 3d77db1 single commit
• 517b763 Removes flax.optim.dynamic_scale.
• 28291a8 Update minimal JAX version
• 890a391 Merge pull request #2389 from mattjj:skip-buggy-test
• 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)