Haiku 0.0.16

What's Changed

• Fix inspect.stack() error caused by no_flax shim module. 58b4130.
• Various changes to remain compatible with the latest JAX release.

Full Changelog: v0.0.15...v0.0.16