Bump dm-haiku from 0.0.2 to 0.0.9

[dependabot]

Bumps dm-haiku from 0.0.2 to 0.0.9.

Release notes

Sourced from dm-haiku's releases.

Haiku 0.0.9

What's Changed

• Support vmap where in_axes is a list rather than a tuple in google-deepmind/dm-haiku@307cf7d
• Pass pmap axis specs optionally to make_model_info in google-deepmind/dm-haiku@d0ba451
• Remove use of jax_experimental_name_stack flag in google-deepmind/dm-haiku@dbc0b1f
• Add param_axis argument to RMSNorm to allow setting scale param shape in google-deepmind/dm-haiku@a4998a0
• Add documentation and error messages for w_init and w_init_scale to avoid confusion in deepmind/dm-haiku#541
• Fix hk.while_loop carrying state when reserving variable sizes of rng keys. by @​copybara-service in deepmind/dm-haiku#551
• Add ensemble example to hk.lift documentation. by @​copybara-service in deepmind/dm-haiku#556

Full Changelog: google-deepmind/dm-haiku@v0.0.8...v0.0.9

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.

... (truncated)

Commits

• 889bb06 Cut Haiku 0.0.9 release.
• a4998a0 Add param_axis argument to RMSNorm to allow setting scale param shape.
• dbc0b1f Remove use of jax_experimental_name_stack flag in preparation for its deletion
• 32781b9 Skip some close over tests when running internal presubmits.
• 9c07b2b Add ensemble example to hk.lift documentation.
• 985be73 Improve split_rng hint error by checking if split_rng is set to True + init.
• d0ba451 Pass pmap axis specs optionally to make_model_info.
• a191d4e Internal change.
• 1924f75 import collections.abc instead of just collections
• 6e602af Fix hk.while_loop carrying state when reserving variable sizes of rng keys.
• 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)