What's Changed
- Migrate references to
jax.lib.xla_extensiontojax.extend. by @copybara-service[bot] in #403 - Add
assert_tree_shape, which is likeassert_shape, but for all leaves in a tree. by @copybara-service[bot] in #404 - Use
format_shape_matcherwhen constructing theassert_tree_shapeerror message. by @copybara-service[bot] in #405 - Migrate to pyproject.toml by @copybara-service[bot] in #406
- release v0.1.91 by @copybara-service[bot] in #407
Full Changelog: v0.1.90...v0.1.91