Bump dm-haiku from 0.0.2 to 0.0.6

[dependabot]

Bumps dm-haiku from 0.0.2 to 0.0.6.

Release notes

Sourced from dm-haiku's releases.

Haiku 0.0.6

• Haiku now returns plain nested dicts rather than FlatMap from all APIs.
• hk.vmap(..) now optionally takes split_rng, this argument will be required in the next version of Haiku.
• hk.LayerNorm now accepts param_axis in the constructor, this argument will be required in the next version of Haiku.
• hk.get_channel_index(..) was added.
• hk.experimental.lift_with_state(..) was added.
• hk.experimental.layer_stack(..) was added.
• hk.DepthwiseConv{1,3}D were added.
• hk.BatchNorm now supports sequences in cross_replica_axis.
• hk.experimental.check_jax_usage() makes Haiku check that JAX control flow/transforms are used correctly.
• hk.dynamic_unroll(..) now supports return_all_states.
• hk.cond(..) supports N operands.
• hk.experimental.module_auto_repr(False) can be used to speed up init.
• hk.data_structures.merge(..) now supports check_duplicates.
• TruncatedNormal initialiser now supports complex dtypes.
• transform(jit(f)) now provides a nice error message.
• hk.multinomial(..) now uses jax.random.categorical.
• Added hk.mixed_precision.{current,get}_policy(..) for introspection.
• Mixed precision policies now support reloaded modules.

Haiku 0.0.5

• Added support for mixed precision training (dba1fd9) via jmp
• Added hk.with_empty_state(..).
• Added hk.multi_transform(..) (#137), supporting transforming multiple functions that share parameters.
• Added hk.data_structures.is_subset(..) to test whether parameters are a subset of another.
• Minimum Python version is now 3.7.
• Multiple changes in preparation for a future version of Haiku changing to plain dicts.
• hk.next_rng_keys(..) now returns a stacked array rather than a collection.
• hk.MultiHeadAttention now supports distinct sequence lengths in query and key/value.
• hk.LayerNorm now optionally supports faster (but less stable) variance computation.
• hk.nets.MLP now has an output_shape property.
• hk.nets.ResNet now supports changing strides.
• UnexpectedTracerError inside a Haiku transform now has a more useful error message.
• hk.{lift,custom_creator,custom_getter} are no longer experimental.
• Haiku now supports JAX's pluggable RNGs.
• We have made multiple improvements to our docs an error messages.

Any many other small fixes and improvements.

Haiku 0.0.4

Changelog:

• (Important Fix) Fixed strides in basic block (300e6a40be3).
• Added map, partition_n and traverse to data_structures.
• Added "build your own Haiku" to the docs.
• Added summarise utility to Haiku.
• Added visualisation section to docs.
• Added precision arg to Linear, Conv and ConvTranspose.
• Added RMSNorm.
• Added module_name and name to GetterContext.

... (truncated)

Commits

• f25eb03 Cut Haiku 0.0.6 release.
• 8d8555d Rename safe_spec to check_duplicates and handle array vs. non-array case.
• 0812512 Add safe merging of parameter structures which checks for conflicting shape o...
• 480c56c Add a global and local way to disable generating automatic repr.
• 0f87a53 Add support for N operand new cond.
• 232cc14 Adding return_all_states options to dynamic_unroll.
• 1254b96 Update Generator[Batch, None, None] types to Iterator[Batch] in examples for ...
• b4d4507 Work around a pytype error.
• 95c8bab Refine valid RNG check to use the expected PRNGimpl key_shape.
• 7bbecdc Add hk.experimental.check_jax_usage().
• 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)

[dependabot]

Superseded by #351.