Bump dm-haiku from 0.0.2 to 0.0.8

[dependabot]

Bumps dm-haiku from 0.0.2 to 0.0.8.

Release notes

Sourced from dm-haiku's releases.

Haiku 0.0.8

• Added experimental.force_name.
• Added ability to simulate a method name in experimental.name_scope.
• Added a config option for PRNG key block size.
• Added unroll parameter to dynamic_unroll.
• Remove use of deprecated jax.tree_* functions.
• Many improvements to our examples.
• Improve error messages in vmap.
• Support jax_experimental_name_stack in jaxpr_info.
• transform_and_run now supports a map on PRNG keys.
• remat now uses the new JAX remat implementation.
• Scale parameter is now optional in RMSNorm.

Haiku 0.0.7

• Bug fix: modules with leading zeros (e.g. linear_007) are now correctly handled. 7632affafad
• Breaking change: hk.vmap(..) now requires split_rng to be passed.
• Breaking change: hk.jit was removed from the public API.
• Feature: we always add profiler name scopes to Haiku modules with the latest version of JAX.
• Added a tutorial on parameter sharing.
• Added hk.ModuleProtocol and hk.SupportsCall.
• Added cross_replica_axis to VectorQuantiser.
• Added allow_reuse argument to hk.lift.
• Added fan_in_axes to VarianceScaling initialiser.
• Added hk.custom_setter(..) to intercept hk.set_state(..).
• Added hk.Deferred.
• Added hk.experimental.transparent_lift(..) and hk.experimental.transparent_lift_with_state(..).
• Added hk.experimental.fast_eval_shape(..).
• Added hk.experimental.current_name().
• Added hk.experimental.DO_NOT_STORE. 2a6c03406f
• Added config APIs.
• Added support for new jax.named_call implementation.
• The HAIKU_FLATMAPPING env var is no longer used.
• hk.dropout(..) now supports dynamic rate.
• hk.without_apply_rng(..) now supports multi transformed functions.

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.

... (truncated)

Commits

• 606e18d Cut Haiku 0.0.8 release.
• e8914b6 Make block size of PRNG key reservation configurable.
• bafdbb5 Make this code compatible with Python 3.10.
• 7f215c1 Switch to public remat API now that new implementation is default.
• 15d55f5 Remove jax.core.UnitVar usage (there's no more units!)
• 93a33bc Improve the error message by saying what it is in addition to what it should be.
• c1c5fad Fix invocation of jax2tf.convert
• 93e174c Bump nbsphinx to fix RTD PDF build.
• 8759cfb Add an option to simulate method names in name_scope.
• b0a7d50 Add hk.experimental.force_name(..).
• 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 #397.