Update flax requirement from <0.4,>=0.3.5 to >=0.3.5,<0.5

Created by: dependabot[bot]

Updates the requirements on flax to permit the latest version.

Release notes

Sourced from flax's releases.

Version 0.4.0

What's Changed

• Add PReLU Activation by @​isaaccorley in google/flax#1570
• Fix GroupNorm type hint for param num_groups. by @​lkhphuc in google/flax#1657
• Add named_call overrides to docs by @​jheek in google/flax#1649
• mission statement by @​jheek in google/flax#1668
• Improves Flax Modules for RTD by @​marcvanzee in google/flax#1416
• Add clarifying docstring for 'size' argument to prefetch_to_device's by @​avital in google/flax#1574
• Add circular padding to flax.linen.Conv and flax.linen.ConvTranspose by @​sgrigory in google/flax#1661
• Fix child scope rng reuse. by @​jheek in google/flax#1692
• Numerically stable weight norm by @​jheek in google/flax#1693
• Remove cyclic refs from scope by @​jheek in google/flax#1696
• Add unroll to jax_utils.scan_in_dim by @​ptigwe in google/flax#1691
• Removes rng arguments from Dropout's __call__. by @​copybara-service in google/flax#1689
• Add error for empty scopes. by @​jheek in google/flax#1698
• correct axis resolution in case of repeated axis in the logica axis r… by @​ultrons in google/flax#1703
• Fix lost mutation bug in transforms on nested scopes. by @​levskaya in google/flax#1716
• Expose put_variable function to Module. by @​levskaya in google/flax#1710
• add eq and hash for scopes by @​jheek in google/flax#1720
• Fixes a bug in DenseGeneral. by @​copybara-service in google/flax#1722
• Add param_dtype argument to linen Modules by @​jheek in google/flax#1739
• Implement custom vjp by @​jheek in google/flax#1738
• Handle setup with transformed methods taking submodules of self. by @​levskaya in google/flax#1745
• validate RNG key shape against jax's default by @​copybara-service in google/flax#1780
• Adds optax update guide. by @​andsteing in google/flax#1774
• Implement LazyRNG by @​jheek in google/flax#1723
• make params_with_axes() work when params_axes is not mutable by @​copybara-service in google/flax#1811
• Updates the ensembling HOWTO to Optax. by @​andsteing in google/flax#1806
• Adds prominent scenic link to examples/README.md by @​copybara-service in google/flax#1809
• Removes PixelCNN++ example. @​copybara-service in google/flax#1819
• Add support for non-float32 normalization for linen normalization layers by @​jheek in google/flax#1804
• Make Filter a Collection instead of a Container by @​NeilGirdhar in google/flax#1815
• Removes deprecated API from RTD by @​marcvanzee in google/flax#1824

New Contributors

• @​isaaccorley made their first contribution in google/flax#1570
• @​lkhphuc made their first contribution in google/flax#1657
• @​sgrigory made their first contribution in google/flax#1661
• @​ptigwe made their first contribution in google/flax#1691
• @​ultrons made their first contribution in google/flax#1703
• @​dependabot made their first contribution in google/flax#1749
• @​NeilGirdhar made their first contribution in google/flax#1699
• @​saeta made their first contribution in google/flax#1784
• @​melissatan made their first contribution in google/flax#1793

Full Changelog: https://github.com/google/flax/compare/v0.3.6...v0.4.0

Changelog

Sourced from flax's changelog.

0.4.0

Breaking changes:

• flax.deprecated.nn is removed. Please pin to flax==0.3.6 if you are still using it.
• PixelCNN++ example is removed. It was not working well on TPU.
• linen Normalization layers no longer downcast double and complex floats tofloat32 when computing the mean and variance.

New features:

• Added flax.linen.custom_vjp for custom derivatives inside a Module.
• Add param_dtype attribute to standard Linen Modules for specifying parameter dtypes.

0.3.6

Breaking changes:

• Move flax.nn to flax.deprecated.nn.

New features:

• Add experimental checkpoint policy argument. See flax.linen.checkpoint
• Add lifted versions of jvp and vjp.
• Add lifted transformation for mapping variables. See flax.linen.map_variables.

0.3.5

Breaking changes:

• You can no longer pass an int as the kernel_size for a `flax.linen.Conv. Instead a type error is raised stating that a tuple/list should be provided. Stride and dilation arguments do support broadcasting a single int value now because this is not ambigious when the kernel rank is known.
• flax.linen.enable_named_call and flax.linen.disable_named_call now work anywhere instead of only affecting Modules constructed after the enable/disable call. Additionally, there is now flax.linen.override_named_call that provided a context manager to locally disable/enable named_call.
• NamedTuples are no longer converted to tuples on assignment to a linen.Module.

New features:

• Flax internal stack frames are now removed from exception state traces.
• Added flax.linen.nowrap to decorate method that should not be transformed because they are stateful.
• Flax no longer uses implicit rank broadcasting. Thus, you can now use Flax with --jax_numpy_rank_promotion=raise.

Bugfixes:

• linen Modules and dataclasses made with flax.struct.dataclass or flax.struct.PyTreeNode are now correctly recognized as dataclasses by static analysis tools like PyLance. Autocomplete of constructors has been verified to work with VSCode.
• Fixed a bug in FrozenDict which didn't allow copying dicts with reserved names.
• Fix the serialization of named tuples. Tuple fields are no longer stored in the state dict and the named tuple class is no longer recreated (bug).
• Mixed precision training with float16 now works correctly with the attention layers.
• auto-generated linen Module __hash__, __eq__, __repr__ no longer fail by default on non-init attributes.

... (truncated)

Commits

• 78d0566 Merge pull request #1825 from marcvanzee:pipy
• 695a880 Merge branch 'main' of https://www.github.com/google/flax into pipy
• 5de87fa Bump Flax to v0.4.0
• a6e7744 Merge pull request #1824 from marcvanzee:rtd-deprecated
• a7cf071 Updates Linen README
• 1c4fe91 Updates core notebook to Linen
• 12a1159 Removes deprecated API from docs
• aca40ad Actually removes the deprecated folder.
• 948ee26 Removes flax.deprecated.nn.
• 7288ada Merge pull request #1815 from NeilGirdhar:container
• 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)