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.