dm-haiku

Add support for init=None in apply.

Make hk.switch support multiple operands just like jax.lax.switch.

Improve error message for invalid rng when state={}.

Minor improvements to layer_stack.

Avoid crashes when constructing signatures