{"id":13806579,"url":"https://github.com/google-deepmind/chex","last_synced_at":"2025-12-11T21:03:46.108Z","repository":{"id":37483722,"uuid":"285530253","full_name":"google-deepmind/chex","owner":"google-deepmind","description":null,"archived":false,"fork":false,"pushed_at":"2025-05-01T13:48:58.000Z","size":594,"stargazers_count":848,"open_issues_count":50,"forks_count":51,"subscribers_count":15,"default_branch":"master","last_synced_at":"2025-05-08T04:02:07.295Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"https://chex.readthedocs.io","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/google-deepmind.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2020-08-06T09:32:36.000Z","updated_at":"2025-05-07T09:09:43.000Z","dependencies_parsed_at":"2023-09-07T20:34:49.970Z","dependency_job_id":"2e8a021f-2ee6-4092-a03f-4451eb2ec280","html_url":"https://github.com/google-deepmind/chex","commit_stats":null,"previous_names":["google-deepmind/chex","deepmind/chex"],"tags_count":23,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fchex","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fchex/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fchex/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-deepmind%2Fchex/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/google-deepmind","download_url":"https://codeload.github.com/google-deepmind/chex/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254159650,"owners_count":22024563,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-08-04T01:01:13.522Z","updated_at":"2025-12-11T21:03:40.861Z","avatar_url":"https://github.com/google-deepmind.png","language":"Python","funding_links":[],"categories":["Libraries","Python"],"sub_categories":[],"readme":"# Chex\n\n![CI status](https://github.com/deepmind/chex/workflows/ci/badge.svg)\n![docs](https://readthedocs.org/projects/chex/badge/?version=latest)\n![pypi](https://img.shields.io/pypi/v/chex)\n\nChex is a library of utilities for helping to write reliable JAX code.\n\nThis includes utils to help:\n\n* Instrument your code (e.g. assertions, warnings)\n* Debug (e.g. transforming `pmaps` in `vmaps` within a context manager).\n* Test JAX code across many `variants` (e.g. jitted vs non-jitted).\n\n## Installation\n\nYou can install the latest released version of Chex from PyPI via:\n\n```sh\npip install chex\n```\n\nor you can install the latest development version from GitHub:\n\n```sh\npip install git+https://github.com/deepmind/chex.git\n```\n\n## Modules Overview\n\n### Dataclass ([dataclass.py](https://github.com/deepmind/chex/blob/master/chex/_src/dataclass.py))\n\nDataclasses are a popular construct introduced by Python 3.7 to allow to\neasily specify typed data structures with minimal boilerplate code. They are\nnot, however, compatible with JAX and\n[dm-tree](https://github.com/deepmind/tree) out of the box.\n\nIn Chex we provide a JAX-friendly dataclass implementation reusing python [dataclasses](https://docs.python.org/3/library/dataclasses.html#module-dataclasses).\n\nChex implementation of `dataclass` registers dataclasses as internal [_PyTree_\nnodes](https://jax.readthedocs.io/en/latest/pytrees.html) to ensure\ncompatibility with JAX data structures.\n\nIn addition, we provide a class wrapper that exposes dataclasses as\n`collections.Mapping` descendants which allows to process them\n(e.g. (un-)flatten) in `dm-tree` methods as usual Python dictionaries.\nSee [`@mappable_dataclass`](https://github.com/deepmind/chex/blob/master/chex/_src/dataclass.py#L27)\ndocstring for more details.\n\nExample:\n\n```python\n@chex.dataclass\nclass Parameters:\n  x: chex.ArrayDevice\n  y: chex.ArrayDevice\n\nparameters = Parameters(\n    x=jnp.ones((2, 2)),\n    y=jnp.ones((1, 2)),\n)\n\n# Dataclasses can be treated as JAX pytrees\njax.tree_util.tree_map(lambda x: 2.0 * x, parameters)\n\n# and as mappings by dm-tree\ntree.flatten(parameters)\n```\n\n**NOTE**: Unlike standard Python 3.7 dataclasses, Chex\ndataclasses cannot be constructed using positional arguments. They support\nconstruction arguments provided in the same format as the Python dict\nconstructor. Dataclasses can be converted to tuples with the `from_tuple` and\n`to_tuple` methods if necessary.\n\n```python\nparameters = Parameters(\n    jnp.ones((2, 2)),\n    jnp.ones((1, 2)),\n)\n# ValueError: Mappable dataclass constructor doesn't support positional args.\n```\n\n### Assertions ([asserts.py](https://github.com/deepmind/chex/blob/master/chex/_src/asserts.py))\n\nOne limitation of PyType annotations for JAX is that they do not support the\nspecification of `DeviceArray` ranks, shapes or dtypes. Chex includes a number\nof functions that allow flexible and concise specification of these properties.\n\nE.g. suppose you want to ensure that all tensors `t1`, `t2`, `t3` have the same\nshape, and that tensors `t4`, `t5` have rank `2` and (`3` or `4`), respectively.\n\n```python\nchex.assert_equal_shape([t1, t2, t3])\nchex.assert_rank([t4, t5], [2, {3, 4}])\n```\n\nMore examples:\n\n```python\nfrom chex import assert_shape, assert_rank, ...\n\nassert_shape(x, (2, 3))                # x has shape (2, 3)\nassert_shape([x, y], [(), (2,3)])      # x is scalar and y has shape (2, 3)\n\nassert_rank(x, 0)                      # x is scalar\nassert_rank([x, y], [0, 2])            # x is scalar and y is a rank-2 array\nassert_rank([x, y], {0, 2})            # x and y are scalar OR rank-2 arrays\n\nassert_type(x, int)                    # x has type `int` (x can be an array)\nassert_type([x, y], [int, float])      # x has type `int` and y has type `float`\n\nassert_equal_shape([x, y, z])          # x, y, and z have equal shapes\n\nassert_trees_all_close(tree_x, tree_y) # values and structure of trees match\nassert_tree_all_finite(tree_x)         # all tree_x leaves are finite\n\nassert_devices_available(2, 'gpu')     # 2 GPUs available\nassert_tpu_available()                 # at least 1 TPU available\n\nassert_numerical_grads(f, (x, y), j)   # f^{(j)}(x, y) matches numerical grads\n```\n\nSee `asserts.py`\n[documentation](https://chex.readthedocs.io/en/latest/api.html#assertions) to\nfind all supported assertions.\n\nIf you cannot find a specific assertion, please consider making a pull request\nor openning an issue on\n[the bug tracker](https://github.com/deepmind/chex/issues).\n\n#### Optional Arguments\n\nAll chex assertions support the following optional kwargs for manipulating the\nemitted exception messages:\n\n* `custom_message`: A string to include into the emitted exception messages.\n* `include_default_message`: Whether to include the default Chex message into\n  the emitted exception messages.\n* `exception_type`: An exception type to use. `AssertionError` by default.\n\nFor example, the following code:\n\n```python\ndataset = load_dataset()\nparams = init_params()\nfor i in range(num_steps):\n  params = update_params(params, dataset.sample())\n  chex.assert_tree_all_finite(params,\n                              custom_message=f'Failed at iteration {i}.',\n                              exception_type=ValueError)\n```\n\nwill raise a `ValueError` that includes a step number when `params` get polluted\nwith `NaNs` or `None`s.\n\n#### Static and Value (aka *Runtime*) Assertions\n\nChex divides all assertions into 2 classes: ***static*** and ***value***\nassertions.\n\n1.  ***static*** assertions use anything except concrete values of tensors.\n    Examples: `assert_shape`, `assert_trees_all_equal_dtypes`,\n    `assert_max_traces`.\n\n2.  ***value*** assertions require access to tensor values, which are not\n    available during JAX tracing (see\n    [HowJAX primitives work](https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primitives_work.html)),\n    thus such assertion need special treatment in a *jitted* code.\n\nTo enable value assertions in a jitted function, it can be decorated with\n`chex.chexify()` wrapper. Example:\n\n```python\n  @chex.chexify\n  @jax.jit\n  def logp1_abs_safe(x: chex.Array) -\u003e chex.Array:\n    chex.assert_tree_all_finite(x)\n    return jnp.log(jnp.abs(x) + 1)\n\n  logp1_abs_safe(jnp.ones(2))  # OK\n  logp1_abs_safe(jnp.array([jnp.nan, 3]))  # FAILS (in async mode)\n\n  # The error will be raised either at the next line OR at the next\n  # `logp1_abs_safe` call. See the docs for more detain on async mode.\n  logp1_abs_safe.wait_checks()  # Wait for the (async) computation to complete.\n```\n\nSee\n[this docstring](https://chex.readthedocs.io/en/latest/api.html#chex.chexify)\nfor more detail on `chex.chexify()`.\n\n#### JAX Tracing Assertions\n\nJAX re-traces JIT'ted function every time the structure of passed arguments\nchanges. Often this behavior is inadvertent and leads to a significant\nperformance drop which is hard to debug. [@chex.assert_max_traces](https://github.com/deepmind/chex/blob/master/chex/_src/asserts.py#L44)\ndecorator asserts that the function is not re-traced more than `n` times during\nprogram execution.\n\nGlobal trace counter can be cleared by calling\n`chex.clear_trace_counter()`. This function be used to isolate unittests relying\non `@chex.assert_max_traces`.\n\nExamples:\n\n```python\n  @jax.jit\n  @chex.assert_max_traces(n=1)\n  def fn_sum_jitted(x, y):\n    return x + y\n\n  fn_sum_jitted(jnp.zeros(3), jnp.zeros(3))  # tracing for the 1st time - OK\n  fn_sum_jitted(jnp.zeros([6, 7]), jnp.zeros([6, 7]))  # AssertionError!\n```\n\nCan be used with `jax.pmap()` as well:\n\n```python\n  def fn_sub(x, y):\n    return x - y\n\n  fn_sub_pmapped = jax.pmap(chex.assert_max_traces(fn_sub, n=10))\n```\n\nSee\n[HowJAX primitives work](https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primitives_work.html)\nsection for more information about tracing.\n\n### Warnings ([warnigns.py](https://github.com/deepmind/chex/blob/master/chex/_src/warnings.py))\n\nIn addition to hard assertions Chex also offers utilities to add common\nwarnings, such as specific types of deprecation warnings.\n\n### Test variants ([variants.py](https://github.com/deepmind/chex/blob/master/chex/_src/variants.py))\n\nJAX relies extensively on code transformation and compilation, meaning that it\ncan be hard to ensure that code is properly tested. For instance, just testing a\npython function using JAX code will not cover the actual code path that is\nexecuted when jitted, and that path will also differ whether the code is jitted\nfor CPU, GPU, or TPU. This has been a source of obscure and hard to catch bugs\nwhere XLA changes would lead to undesirable behaviours that however only\nmanifest in one specific code transformation.\n\nVariants make it easy to ensure that unit tests cover different ‘variations’ of\na function, by providing a simple decorator that can be used to repeat any test\nunder all (or a subset) of the relevant code transformations.\n\nE.g. suppose you want to test the output of a function `fn` with or without jit.\nYou can use `chex.variants` to run the test with both the jitted and non-jitted\nversion of the function by simply decorating a test method with\n`@chex.variants`, and then using `self.variant(fn)` in place of `fn` in the body\nof the test.\n\n```python\ndef fn(x, y):\n  return x + y\n...\n\nclass ExampleTest(chex.TestCase):\n\n  @chex.variants(with_jit=True, without_jit=True)\n  def test(self):\n    var_fn = self.variant(fn)\n    self.assertEqual(fn(1, 2), 3)\n    self.assertEqual(var_fn(1, 2), fn(1, 2))\n```\n\nIf you define the function in the test method, you may also use `self.variant`\nas a decorator in the function definition. For example:\n\n```python\nclass ExampleTest(chex.TestCase):\n\n  @chex.variants(with_jit=True, without_jit=True)\n  def test(self):\n    @self.variant\n    def var_fn(x, y):\n       return x + y\n\n    self.assertEqual(var_fn(1, 2), 3)\n```\n\nExample of parameterized test:\n\n```python\nfrom absl.testing import parameterized\n\n# Could also be:\n#  `class ExampleParameterizedTest(chex.TestCase, parameterized.TestCase):`\n#  `class ExampleParameterizedTest(chex.TestCase):`\nclass ExampleParameterizedTest(parameterized.TestCase):\n\n  @chex.variants(with_jit=True, without_jit=True)\n  @parameterized.named_parameters(\n      ('case_positive', 1, 2, 3),\n      ('case_negative', -1, -2, -3),\n  )\n  def test(self, arg_1, arg_2, expected):\n    @self.variant\n    def var_fn(x, y):\n       return x + y\n\n    self.assertEqual(var_fn(arg_1, arg_2), expected)\n```\n\nChex currently supports the following variants:\n\n* `with_jit` -- applies `jax.jit()` transformation to the function.\n* `without_jit` -- uses the function as is, i.e. identity transformation.\n* `with_device` -- places all arguments (except specified in `ignore_argnums`\n   argument) into device memory before applying the function.\n* `without_device` -- places all arguments in RAM before applying the function.\n* `with_pmap` -- applies `jax.pmap()` transformation to the function (see notes below).\n\nSee documentation in [variants.py](https://github.com/deepmind/chex/blob/master/chex/_src/variants.py) for more details on the supported variants.\nMore examples can be found in [variants_test.py](https://github.com/deepmind/chex/blob/master/chex/_src/variants_test.py).\n\n### Variants notes\n\n* Test classes that use `@chex.variants` must inherit from\n`chex.TestCase` (or any other base class that unrolls tests generators\nwithin `TestCase`, e.g. `absl.testing.parameterized.TestCase`).\n\n* **[`jax.vmap`]** All variants can be applied to a vmapped function;\nplease see an example in [variants_test.py](https://github.com/deepmind/chex/blob/master/chex/_src/variants_test.py) (`test_vmapped_fn_named_params` and\n`test_pmap_vmapped_fn`).\n\n* **[`@chex.all_variants`]** You can get all supported variants\nby using the decorator `@chex.all_variants`.\n\n* **[`with_pmap` variant]** `jax.pmap(fn)`\n([doc](https://jax.readthedocs.io/en/latest/jax.html#jax.pmap)) performs\nparallel map of `fn` onto multiple devices. Since most tests run in a\nsingle-device environment (i.e. having access to a single CPU or GPU), in which\ncase `jax.pmap` is a functional equivalent to `jax.jit`, ` with_pmap` variant is\nskipped by default (although it works fine with a single device). Below we\ndescribe  a way to properly test `fn` if it is supposed to be used in\nmulti-device environments (TPUs or multiple CPUs/GPUs). To disable skipping\n`with_pmap` variants in case of a single device, add\n`--chex_skip_pmap_variant_if_single_device=false` to your test command.\n\n### Fakes ([fake.py](https://github.com/deepmind/chex/blob/master/chex/_src/fake.py))\n\nDebugging in JAX is made more difficult by code transformations such as `jit`\nand `pmap`, which introduce optimizations that make code hard to inspect and\ntrace. It can also be difficult to disable those transformations during\ndebugging as they can be called at several places in the underlying\ncode. Chex provides tools to globally replace `jax.jit` with a no-op\ntransformation and `jax.pmap` with a (non-parallel) `jax.vmap`, in order to more\neasily debug code in a single-device context.\n\nFor example, you can use Chex to fake `pmap` and have it replaced with a `vmap`.\nThis can be achieved by wrapping your code with a context manager:\n\n```python\nwith chex.fake_pmap():\n  @jax.pmap\n  def fn(inputs):\n    ...\n\n  # Function will be vmapped over inputs\n  fn(inputs)\n```\n\nThe same functionality can also be invoked with `start` and `stop`:\n\n```python\nfake_pmap = chex.fake_pmap()\nfake_pmap.start()\n... your jax code ...\nfake_pmap.stop()\n```\n\nIn addition, you can fake a real multi-device test environment with a\nmulti-threaded CPU. See section **Faking multi-device test environments** for\nmore details.\n\nSee documentation in [fake.py](https://github.com/deepmind/chex/blob/master/chex/_src/fake.py) and examples in [fake_test.py](https://github.com/deepmind/chex/blob/master/chex/_src/fake_test.py) for more details.\n\n## Faking multi-device test environments\n\nIn situations where you do not have easy access to multiple devices, you can\nstill test parallel computation using single-device multi-threading.\n\nIn particular, one can force XLA to use a single CPU's threads as separate\ndevices, i.e. to fake a real multi-device environment with a multi-threaded one.\nThese two options are theoretically equivalent from XLA perspective because they\nexpose the same interface and use identical abstractions.\n\nChex has a flag `chex_n_cpu_devices` that specifies a number of CPU threads to\nuse as XLA devices.\n\nTo set up a multi-threaded XLA environment for `absl` tests, define\n`setUpModule` function in your test module:\n\n```python\ndef setUpModule():\n  chex.set_n_cpu_devices()\n```\n\nNow you can launch your test with `python test.py --chex_n_cpu_devices=N` to run\nit in multi-device regime. Note that **all** tests within a module will have an\naccess to `N` devices.\n\nMore examples can be found in [variants_test.py](https://github.com/deepmind/chex/blob/master/chex/_src/variants_test.py), [fake_test.py](https://github.com/deepmind/chex/blob/master/chex/_src/fake_test.py) and [fake_set_n_cpu_devices_test.py](https://github.com/deepmind/chex/blob/master/chex/_src/fake_set_n_cpu_devices_test.py).\n\n### Using named dimension sizes.\n\nChex comes with a small utility that allows you to package a collection of\ndimension sizes into a single object. The basic idea is:\n\n```python\ndims = chex.Dimensions(B=batch_size, T=sequence_len, E=embedding_dim)\n...\nchex.assert_shape(arr, dims['BTE'])\n```\n\nString lookups are translated integer tuples. For instance, let's say\n`batch_size == 3`, `sequence_len = 5` and `embedding_dim = 7`, then\n\n```python\ndims['BTE'] == (3, 5, 7)\ndims['B'] == (3,)\ndims['TTBEE'] == (5, 5, 3, 7, 7)\n...\n```\n\nYou can also assign dimension sizes dynamically as follows:\n\n```python\ndims['XY'] = some_matrix.shape\ndims.Z = 13\n```\n\nFor more examples, see [chex.Dimensions](https://chex.readthedocs.io/en/latest/api.html#chex.Dimensions)\ndocumentation.\n\n## Citing Chex\n\nThis repository is part of the [DeepMind JAX Ecosystem], to cite Chex please use\nthe [DeepMind JAX Ecosystem citation].\n\n[DeepMind JAX Ecosystem]: https://deepmind.com/blog/article/using-jax-to-accelerate-our-research \"DeepMind JAX Ecosystem\"\n[DeepMind JAX Ecosystem citation]: https://github.com/deepmind/jax/blob/main/deepmind2020jax.txt \"Citation\"\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Fchex","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fgoogle-deepmind%2Fchex","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-deepmind%2Fchex/lists"}