Haiku 0.0.17

What's Changed

• 9396aea Exposed spmd_axis_name in hk.vmap.
• f69e304 Improve static_unroll VJP performance by slicing via jnp.unstack.
• Misc changes to remain compatible with newer versions of JAX.

Full Changelog: v0.0.16...v0.0.17