Chex

It is the latest version compatible with Python 3.6. See https://github.com/deepmind/optax/issues/222 for more details.

Changes:

* Add support for static_broadcasted_argnums to fake_pmap;
* Allows sets of alternatives and ellipsis in assert_shape;
* Format variant test names to use only underscores and lowercase letters;
* Fix incorrect type annotation in asserts.py;
* Fix dataclass (un-)flatten functions;
* Add more tests for dataclasses;
* Raise ValueError when no variants are selected;
* Exclude chex' internal frames from AssertionError tracebacks;
* Add '[Chex] ' prefix to AssertionError messages;
* Include path to leaves that failed the equality check in `assert_tree_all_close`;
* Clean up asserts.py;
* Asserts which only make sense on >1 tree now demand this (can result in breakages in the existing code).

- Add is\_divisible assertion. [\49](https://github.com/deepmind/chex/pull/49) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Create py.typed [\46](https://github.com/deepmind/chex/pull/46) ([graingert](https://github.com/graingert))
- Remove "\ Lint as: python3" directives. [\45](https://github.com/deepmind/chex/pull/45) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Pass arguments to jax.jit other than the function as kwargs instead of positional arguments. [\43](https://github.com/deepmind/chex/pull/43) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- \[JAX\] Use a hand-written Python wrapper for CompiledFunction. [\42](https://github.com/deepmind/chex/pull/42) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Ensure LICENSE files are an exact copy of Apache LICENSE-2.0.txt \[0\]. [\41](https://github.com/deepmind/chex/pull/41) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix assert\_max\_traces usage checks. [\40](https://github.com/deepmind/chex/pull/40) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix tree path formatting. [\39](https://github.com/deepmind/chex/pull/39) ([copybara-service[bot]](https://github.com/apps/copybara-service))



\* *This Changelog was automatically generated by [github_changelog_generator](https://github.com/github-changelog-generator/github-changelog-generator)*

Changelog

[Full Changelog](https://github.com/deepmind/chex/compare/v0.0.5...HEAD)

**Merged pull requests:**

- Define MANIFEST.in to include requirements.txt in pypi bundle. Increment the version. [\38](https://github.com/deepmind/chex/pull/38) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix PACKAGE\_VER in pypi-publish.yml [\36](https://github.com/deepmind/chex/pull/36) ([copybara-service[bot]](https://github.com/apps/copybara-service))



\* *This Changelog was automatically generated by [github_changelog_generator](https://github.com/github-changelog-generator/github-changelog-generator)*

Changelog

*Note: this is a first GitHub release of Chex. It includes all changes since the repo was created.*

[Full Changelog](https://github.com/deepmind/chex/compare/a0153acadc3609b89d94bc594764b353617164f5...HEAD)

**Closed issues:**

- Chex dataclass throws an exception in Python 3.9 [\10](https://github.com/deepmind/chex/issues/10)
- 'jax.interpreters.xla' has no attribute '\_DeviceArray' for jax \<= 0.2.5 [\9](https://github.com/deepmind/chex/issues/9)
- Chex dataclass defaulting mappable\_dataclass=True [\8](https://github.com/deepmind/chex/issues/8)
- DeprecationWarning for importing toolz [\4](https://github.com/deepmind/chex/issues/4)
- Fake contexts by calling .start\(\) not working [\3](https://github.com/deepmind/chex/issues/3)

**Merged pull requests:**

- Add PyPI release workflow and increment the version. [\35](https://github.com/deepmind/chex/pull/35) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add .pylintrc and run pylint checks in CI workflow. [\34](https://github.com/deepmind/chex/pull/34) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix typo [\33](https://github.com/deepmind/chex/pull/33) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- \[JAX\] Change Chex variants\_test to use jnp.DeviceArray for isinstance\(\) tests. [\32](https://github.com/deepmind/chex/pull/32) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix chex typing annotations [\31](https://github.com/deepmind/chex/pull/31) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Parallelize pytest runs. [\29](https://github.com/deepmind/chex/pull/29) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Fix assert\_max\_traces usage checks. [\28](https://github.com/deepmind/chex/pull/28) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Customize processing of `None`s in tree asserts and add `assert_tree_no_nones`. [\27](https://github.com/deepmind/chex/pull/27) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add option to fake pmap with vmap while still jitting the result. [\25](https://github.com/deepmind/chex/pull/25) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Raise informative error for negative axes when asserting axis dimension. [\24](https://github.com/deepmind/chex/pull/24) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Expose XLA device types in pytypes. [\23](https://github.com/deepmind/chex/pull/23) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add `assert_tree_shape_prefix` assertion. [\21](https://github.com/deepmind/chex/pull/21) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add a safeguard to setup.py to ensure correct package installation. [\20](https://github.com/deepmind/chex/pull/20) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Disable testing against macos-latest to speed up CI. [\19](https://github.com/deepmind/chex/pull/19) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Do not parse abseil flags to avoid side effects. [\18](https://github.com/deepmind/chex/pull/18) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Update version. [\17](https://github.com/deepmind/chex/pull/17) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Run tests under Python 3.9 [\16](https://github.com/deepmind/chex/pull/16) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Correct and clarify docstring for `assert_equal_shape_suffix`. [\14](https://github.com/deepmind/chex/pull/14) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Use type\(\) instead of \_\_class\_\_ in dataclass\_test.py [\13](https://github.com/deepmind/chex/pull/13) ([hbq1](https://github.com/hbq1))
- Include CI status badge in README.md [\12](https://github.com/deepmind/chex/pull/12) ([copybara-service[bot]](https://github.com/apps/copybara-service))
- Add ability to check shapes with wildcards [\5](https://github.com/deepmind/chex/pull/5) ([KristianHolsheimer](https://github.com/KristianHolsheimer))
- Add pytypes. [\2](https://github.com/deepmind/chex/pull/2) ([hbq1](https://github.com/hbq1))
- fix 404 link for Dataclass module header [\1](https://github.com/deepmind/chex/pull/1) ([kykosic](https://github.com/kykosic))



\* *This Changelog was automatically generated by [github_changelog_generator](https://github.com/github-changelog-generator/github-changelog-generator)*