What's Changed
- Switch from
jax.tree_map(deprecated since JAX 0.4.26) to
jax.tree_util.tree_map. - Improvements to parameter overview.
Full Changelog: v0.0.11...v0.0.12
jax.tree_map (deprecated since JAX 0.4.26) tojax.tree_util.tree_map.Full Changelog: v0.0.11...v0.0.12