What's Changed

• Migrate references to jax.lib.xla_extension to jax.extend. by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/403
• Add assert_tree_shape, which is like assert_shape, but for all leaves in a tree. by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/404
• Use format_shape_matcher when constructing the assert_tree_shape error message. by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/405
• Migrate to pyproject.toml by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/406
• release v0.1.91 by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/407

Full Changelog: https://github.com/google-deepmind/chex/compare/v0.1.90...v0.1.91

What's Changed

• Fix source code links to point to master branch. by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/384
• chex: avoid use of deprecated jax.util.unzip2 by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/391
• Update JAX nightly index usage by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/392
• [chex] remove stale try/except import by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/396
• Update restrict_backends after https://github.com/jax-ml/jax/commit/06448864abd6e8187e5b4d9b1ff08ab14fe3b8e0. by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/397
• Use backend_compile_and_load if available, and backend_compile otherwise. by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/398
• Support BF16 in assert_trees_all_close_ulp. by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/399
• Fixed typo. by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/400
• Release v0.1.90 by @copybara-service[bot] in https://github.com/google-deepmind/chex/pull/402

Full Changelog: https://github.com/google-deepmind/chex/compare/v0.1.89...v0.1.90

What's Changed

• Switch to new unittest test loader API by @daskol in https://github.com/google-deepmind/chex/pull/372
• Add test against JAX nightly by @copybara-service in https://github.com/google-deepmind/chex/pull/376
• CI: test Python 3.12 by @copybara-service in https://github.com/google-deepmind/chex/pull/377
• Chex: add Python 3.13 CI tests by @copybara-service in https://github.com/google-deepmind/chex/pull/378
• Fix typo: "assetion" -> "assertion" by @copybara-service in https://github.com/google-deepmind/chex/pull/379
• Maintain compatibility with Python 3.13 by @daskol in https://github.com/google-deepmind/chex/pull/375

New Contributors

• @daskol made their first contribution in https://github.com/google-deepmind/chex/pull/372

Full Changelog: https://github.com/google-deepmind/chex/compare/v0.1.88...v0.1.89

What's Changed

• Fix test_assert_trees_all_equal_prng_keys by @copybara-service in https://github.com/google-deepmind/chex/pull/366
• Avoid relying on jax.lib.xla_bridge._backends, as this is private and will soon be deleted by @copybara-service in https://github.com/google-deepmind/chex/pull/367
• Release 0.1.88. by @copybara-service in https://github.com/google-deepmind/chex/pull/369

Full Changelog: https://github.com/google-deepmind/chex/compare/v0.1.87...v0.1.88

What's Changed

• Add flattening functionality to chex.Dimensions. by @copybara-service in https://github.com/google-deepmind/chex/pull/341
• Fix doc-string in warn_deprecated_function. by @copybara-service in https://github.com/google-deepmind/chex/pull/343
• Add utility to create a deprecated alias. by @copybara-service in https://github.com/google-deepmind/chex/pull/345
• Fix linter's warnings. by @copybara-service in https://github.com/google-deepmind/chex/pull/349
• Test RTD webhooks. by @copybara-service in https://github.com/google-deepmind/chex/pull/350
• Remove a bad async chexify test as the following is not guaranteed: by @copybara-service in https://github.com/google-deepmind/chex/pull/353
• Replace deprecated jax.tree_* functions with jax.tree.* by @copybara-service in https://github.com/google-deepmind/chex/pull/354
• Disable registration of dataclasses defined in main with JAX tree util. by @copybara-service in https://github.com/google-deepmind/chex/pull/356
• Chex: Add CI job for oldest supported JAX version by @copybara-service in https://github.com/google-deepmind/chex/pull/358
• make replacement kwarg in warn_deprecated_function optional by @copybara-service in https://github.com/google-deepmind/chex/pull/362
• increase version number for a pypi release by @copybara-service in https://github.com/google-deepmind/chex/pull/363

Full Changelog: https://github.com/google-deepmind/chex/compare/v0.1.86...v0.1.87

What's Changed

• Avoid use of deprecated device_buffer attriutes of jax.Array by @copybara-service in https://github.com/google-deepmind/chex/pull/325
• Update test.sh by @copybara-service in https://github.com/google-deepmind/chex/pull/331
• Update chex.assert_type to check concrete types instead of just asserting that the type is a floating/integer sub-type. by @copybara-service in https://github.com/google-deepmind/chex/pull/333
• Adjust test tolerances for test_assert_trees_all_close by @copybara-service in https://github.com/google-deepmind/chex/pull/334
• Register dataclass at definition so that unpickling the dataclass object, type, and treedef works in most cases (with one exception). by @copybara-service in https://github.com/google-deepmind/chex/pull/336
• Add warnings in chex. by @copybara-service in https://github.com/google-deepmind/chex/pull/337
• Add Dimensions.size to get flat size from named shape. by @copybara-service in https://github.com/google-deepmind/chex/pull/339
• Increase version number to make a new release. by @copybara-service in https://github.com/google-deepmind/chex/pull/342

Full Changelog: https://github.com/google-deepmind/chex/compare/v0.1.85...v0.1.86

What's Changed

• [JAX] Replace uses of jax.devices("cpu") with jax.local_devices(backend="cpu"). by @copybara-service in https://github.com/google-deepmind/chex/pull/317

Full Changelog: https://github.com/google-deepmind/chex/compare/v0.1.84...v0.1.85

Update the type definitions in pytype.py and fix various warnings.

• Remove deprecated assert_tree_* methods;
• Improve the usability of dataclasses;
• Update dependencies and pytypes in accordance with the latest changes in Jax.

What's Changed

Chex has moved from external dm-tree to the JAX-native jax.tree_util library for operating with PyTrees. This means that Nones will not be treated as separate leaves by chex tree assertions any longer. Note that you can still check for Nones using chex.assert_tree_no_nones().

• Remove usage of dm-tree in asserts. by @copybara-service in https://github.com/deepmind/chex/pull/288
• Add support for kw_only for chex.dataclass. by @jblespiau in https://github.com/deepmind/chex/pull/290
• Deprecate ignore_nones. by @copybara-service in https://github.com/deepmind/chex/pull/292

Full Changelog: https://github.com/deepmind/chex/compare/v0.1.81...v0.1.82

Update requirements.txt to use numpy>=1.25.0.

What's Changed

• Drop support for python 3.8, as per JAX deprecation policy.
• Use new jax.tree_util.register_pytree_with_keys to register chex.Dataclass. by @copybara-service in https://github.com/deepmind/chex/pull/266
• Add size based assertions by @tttc3 in https://github.com/deepmind/chex/pull/265
• Integer sequence support for chex.assert_tree_shape{_prefix, _suffix} by @hylkedonker in https://github.com/deepmind/chex/pull/275
• Resolve unsoundness caught by pytype --strict-none-binding. by @copybara-service in https://github.com/deepmind/chex/pull/276
• Release v0.1.8. by @copybara-service in https://github.com/deepmind/chex/pull/281

New Contributors

• @tttc3 made their first contribution in https://github.com/deepmind/chex/pull/265
• @hylkedonker made their first contribution in https://github.com/deepmind/chex/pull/275

Full Changelog: https://github.com/deepmind/chex/compare/v0.1.7...v0.1.8

What's Changed

• Bump the minimum Jax version to 0.4.6
• feat(ci): bump setup-python version and enable cache by @SauravMaheshkar in https://github.com/deepmind/chex/pull/239
• [README] fix link to @mappable_dataclass by @pfackeldey in https://github.com/deepmind/chex/pull/241
• Add ArrayDeviceTree and ArrayNumpyTree pytypes + improve ArrayTree. by @copybara-service in https://github.com/deepmind/chex/pull/253
• Release v0.1.7 by @copybara-service in https://github.com/deepmind/chex/pull/256

New Contributors

• @SauravMaheshkar made their first contribution in https://github.com/deepmind/chex/pull/239
• @pfackeldey made their first contribution in https://github.com/deepmind/chex/pull/241

Full Changelog: https://github.com/deepmind/chex/compare/v0.1.6...v0.1.7

What's Changed

• Replace jnp.DeviceArray with the new public type jax.Array. by @copybara-service in https://github.com/deepmind/chex/pull/187
• Unpack assertions' kwargs at a higher level for improved readability. by @copybara-service in https://github.com/deepmind/chex/pull/209
• Allow calling @chexify-ed functions from non-main threads. by @copybara-service in https://github.com/deepmind/chex/pull/214
• Add assert_trees_all_equal_shapes_and_dtypes by @copybara-service in https://github.com/deepmind/chex/pull/219
• Implement a new unflatten for chex.dataclass that avoids init while keeping the (un)flattened order unchanged. by @copybara-service in https://github.com/deepmind/chex/pull/222
• Fix examples in the docs & linter warnings. by @copybara-service in https://github.com/deepmind/chex/pull/235
• Replace uses of jax.lib.xla_extension.Device with jax.Device. by @copybara-service in https://github.com/deepmind/chex/pull/237

Full Changelog: https://github.com/deepmind/chex/compare/v0.1.5...v0.1.6

What's Changed

• Add support for value assertions in jitted functions. by @copybara-service in https://github.com/deepmind/chex/pull/178
• [JAX] Avoid private implementation detail _ScalarMeta. by @copybara-service in https://github.com/deepmind/chex/pull/180
• [JAX] Avoid implicit references to jax._src. by @copybara-service in https://github.com/deepmind/chex/pull/181
• Release v0.1.15 by @copybara-service in https://github.com/deepmind/chex/pull/184

Full Changelog: https://github.com/deepmind/chex/compare/v0.1.4...v0.1.5

What's Changed

• Add an InitVar field in the dataclass tests. by @copybara-service in https://github.com/deepmind/chex/pull/161
• Download latest .pylintrc version in tests. by @copybara-service in https://github.com/deepmind/chex/pull/167
• Fix assert_axis_dimension_comparator usages. by @copybara-service in https://github.com/deepmind/chex/pull/168
• Update "jax.tree_util" functions by @copybara-service in https://github.com/deepmind/chex/pull/171
• Use jax.tree_util.tree_map in place of deprecated tree_multimap. by @copybara-service in https://github.com/deepmind/chex/pull/175
• Silence some pytype errors. by @copybara-service in https://github.com/deepmind/chex/pull/174
• Add chex.Dimensions utility for readable shape asserts. by @copybara-service in https://github.com/deepmind/chex/pull/169

Full Changelog: https://github.com/deepmind/chex/compare/v0.1.3...v0.1.4

What's Changed

• Slight helping clarification to clear_trace_counter. by @lucasb-eyer in https://github.com/deepmind/chex/pull/148
• Add new JAX-specific pytypes to chex pytypes. by @copybara-service in https://github.com/deepmind/chex/pull/153
• Remove chex.{C,G,T}puDevice in favour of chex.Device. by @copybara-service in https://github.com/deepmind/chex/pull/154

New Contributors

• @lucasb-eyer made their first contribution in https://github.com/deepmind/chex/pull/148

Full Changelog: https://github.com/deepmind/chex/compare/v0.1.2...v0.1.3

What's Changed

• Support JAX parallel operations in chex.fake_pmap contexts by @copybara-service in https://github.com/deepmind/chex/pull/142
• Remove references to jax.numpy.lax_numpy. by @copybara-service in https://github.com/deepmind/chex/pull/150

Full Changelog: https://github.com/deepmind/chex/compare/v0.1.1...v0.1.2

What's Changed

• Move dataclass registration to init so that it's invoked after deserialization. by @copybara-service in https://github.com/deepmind/chex/pull/111
• Add pytype for jax array's dtypes. by @copybara-service in https://github.com/deepmind/chex/pull/112
• Fix dataclass registration on deserialization. by @copybara-service in https://github.com/deepmind/chex/pull/114
• Fix restrict_backends after jax.xla.backend_compile was moved by @copybara-service in https://github.com/deepmind/chex/pull/116
• Refactor asserts.py and warn users not to rely on asserts_internal's functionality. by @copybara-service in https://github.com/deepmind/chex/pull/117
• Set up ReadTheDoc pages and add a few examples. by @copybara-service in https://github.com/deepmind/chex/pull/118
• Include Sphinx builds into CI tests. by @copybara-service in https://github.com/deepmind/chex/pull/119
• Adds internal functionality by @copybara-service in https://github.com/deepmind/chex/pull/122
• Update Chex citation. by @copybara-service in https://github.com/deepmind/chex/pull/125
• Refactor assertions in preparation for including them into the RTD docs. by @copybara-service in https://github.com/deepmind/chex/pull/126
• Add asserts, variants, and pytypes modules to the RTD docs. by @copybara-service in https://github.com/deepmind/chex/pull/127
• Fix references to collections.abc.Mappable -> collections.abc.Mapping in docs and comments. collections.abc.Mappable does not exist. by @copybara-service in https://github.com/deepmind/chex/pull/129
• Document the rational behing the mappability of chex.dataclasses. by @copybara-service in https://github.com/deepmind/chex/pull/130
• Add 3 new tree assertions: by @copybara-service in https://github.com/deepmind/chex/pull/131
• Add assert_tree_is_sharded for asserting that a tree is sharded across the specified devices. by @copybara-service in https://github.com/deepmind/chex/pull/132
• Add PyTreeDef to pytypes. by @copybara-service in https://github.com/deepmind/chex/pull/134
• Disallow ShardedDeviceArrays in assert_tree_is_on_host and assert_tree_is_on_device. by @copybara-service in https://github.com/deepmind/chex/pull/133
• Bump ipython from 7.16.1 to 7.16.3 in /requirements by @dependabot in https://github.com/deepmind/chex/pull/135
• Remove the old venv directory before testing the package. by @copybara-service in https://github.com/deepmind/chex/pull/138
• Refactor asserts.py in preparation for experimental device assertions. by @copybara-service in https://github.com/deepmind/chex/pull/137
• Fix minor typo in docs. by @copybara-service in https://github.com/deepmind/chex/pull/139
• Improve exception message for assert_tree_shape_prefix. by @copybara-service in https://github.com/deepmind/chex/pull/143
• Release v0.1.1 by @copybara-service in https://github.com/deepmind/chex/pull/146

Full Changelog: https://github.com/deepmind/chex/compare/v0.1.0...v0.1.1

Support for Python 3.6 has been dropped as per JAX deprecation policy. Please upgrade to a supported Python version.

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

Changes since 0.0.8:

• Use rtol=1e-6 in asserts.assert_tree_close;
• Added asserts.assert_trees_all_equal;
• Removed restricted_inheritance option from Chex dataclasses;
• Added dims= option to assert_equal_shape, to check a subset of dims;
• Added test.sh for launching CI tests on a local machine;
• Added support for default exception messages and types to assertions;
• Added support for jnp.bfloat16 to asserts.assert_trees_all_close();
• Added support for static_argnames to variants.with_jit;
• Added a restrict_backends module for constraining the set of backends that a region of code can use;
• Added asserts.assert_trees_all_equal_dtypes assertion;
• Exposed asserts.assert_tree_shape_suffix to the public API;
• Added asserts.assert_tree_shape_suffix to check whether arrays share the same suffix.

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).

Changelog

Full Changelog

Closed issues:

• [REQ] Conda recipe #37

Merged pull requests:

• Release Chex 0.0.7 #50 (copybara-service[bot])
• Add is_divisible assertion. #49 (copybara-service[bot])
• Create py.typed #46 (graingert)
• Remove "# Lint as: python3" directives. #45 (copybara-service[bot])
• Pass arguments to jax.jit other than the function as kwargs instead of positional arguments. #43 (copybara-service[bot])
• [JAX] Use a hand-written Python wrapper for CompiledFunction. #42 (copybara-service[bot])
• Ensure LICENSE files are an exact copy of Apache LICENSE-2.0.txt [0]. #41 (copybara-service[bot])
• Fix @assert_max_traces usage checks. #40 (copybara-service[bot])
• Fix tree path formatting. #39 (copybara-service[bot])

* This Changelog was automatically generated by github_changelog_generator

Changelog

Full Changelog

Merged pull requests:

• Define MANIFEST.in to include requirements.txt in pypi bundle. Increment the version. #38 (copybara-service[bot])
• Fix PACKAGE_VER in pypi-publish.yml #36 (copybara-service[bot])

* This Changelog was automatically generated by github_changelog_generator

Changelog

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

Full Changelog

Closed issues:

• Chex dataclass throws an exception in Python 3.9 #10
• 'jax.interpreters.xla' has no attribute '_DeviceArray' for jax <= 0.2.5 #9
• Chex dataclass defaulting mappable_dataclass=True #8
• DeprecationWarning for importing toolz #4
• Fake contexts by calling .start() not working #3

Merged pull requests:

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

* This Changelog was automatically generated by github_changelog_generator