Chex 0.1.86

What's Changed

• Avoid use of deprecated device_buffer attriutes of jax.Array by @copybara-service in #325
• Update test.sh by @copybara-service in #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 #333
• Adjust test tolerances for test_assert_trees_all_close by @copybara-service in #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 #336
• Add warnings in chex. by @copybara-service in #337
• Add Dimensions.size to get flat size from named shape. by @copybara-service in #339
• Increase version number to make a new release. by @copybara-service in #342

Full Changelog: v0.1.85...v0.1.86

Chex 0.1.85

What's Changed

• [JAX] Replace uses of jax.devices("cpu") with jax.local_devices(backend="cpu"). by @copybara-service in #317

Full Changelog: v0.1.84...v0.1.85

Chex 0.1.84

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

Chex 0.1.83

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

Chex 0.1.82

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 #288
• Add support for kw_only for chex.dataclass. by @jblespiau in #290
• Deprecate ignore_nones. by @copybara-service in #292

Full Changelog: v0.1.81...v0.1.82

Chex 0.1.81

Update requirements.txt to use numpy>=1.25.0.

Chex 0.1.8

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 #266
• Add size based assertions by @tttc3 in #265
• Integer sequence support for chex.assert_tree_shape{_prefix, _suffix} by @hylkedonker in #275
• Resolve unsoundness caught by pytype --strict-none-binding. by @copybara-service in #276
• Release v0.1.8. by @copybara-service in #281

New Contributors

• @tttc3 made their first contribution in #265
• @hylkedonker made their first contribution in #275

Full Changelog: v0.1.7...v0.1.8

Chex 0.1.7

What's Changed

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

New Contributors

• @SauravMaheshkar made their first contribution in #239
• @pfackeldey made their first contribution in #241

Full Changelog: v0.1.6...v0.1.7

Chex 0.1.6

What's Changed

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

Full Changelog: v0.1.5...v0.1.6

Chex 0.1.5

What's Changed

• Add support for value assertions in jitted functions. by @copybara-service in #178
• [JAX] Avoid private implementation detail _ScalarMeta. by @copybara-service in #180
• [JAX] Avoid implicit references to jax._src. by @copybara-service in #181
• Release v0.1.15 by @copybara-service in #184

Full Changelog: v0.1.4...v0.1.5