Package: chex
ChexChex is a library of utilities for helping to write reliable JAX code.This includes utils to help:Instrument your code (e.g. assertions)Debug (e.g. transformingpmapsinvmapswithin a context manager).Test JAX code across manyvariants(e.g. jitted vs non-jitted).InstallationYou can install the latest released version of Chex from PyPI via:pipinstallchexor you can install the latest development version from GitHub:pipinstallgit+https://github.com/deepmind/chex.gitModules OverviewDataclass (dataclass.py)Dataclasses are a popular construct introduced by Python 3.7 to allow to
easily specify typed data structures with minimal boilerplate code. They are
not, however, compatible with JAX anddm-treeout of the box.In Chex we provide a JAX-friendly dataclass implementation reusing pythondataclasses.Chex implementation ofdataclassregisters dataclasses as internalPyTreenodesto ensure
compatibility with JAX data structures.In addition, we provide a class wrapper that exposes dataclasses ascollections.Mappingdescendants which allows to process them
(e.g. (un-)flatten) indm-treemethods as usual Python dictionaries.
See@mappable_dataclassdocstring for more details.Example:@chex.dataclassclassParameters:x:chex.ArrayDevicey:chex.ArrayDeviceparameters=Parameters(x=jnp.ones((2,2)),y=jnp.ones((1,2)),)# Dataclasses can be treated as JAX pytreesjax.tree_util.tree_map(lambdax:2.0*x,parameters)# and as mappings by dm-treetree.flatten(parameters)NOTE: Unlike standard Python 3.7 dataclasses, Chex
dataclasses cannot be constructed using positional arguments. They support
construction arguments provided in the same format as the Python dict
constructor. Dataclasses can be converted to tuples with thefrom_tupleandto_tuplemethods if necessary.parameters=Parameters(jnp.ones((2,2)),jnp.ones((1,2)),)# ValueError: Mappable dataclass constructor doesn't support positional args.Assertions (asserts.py)One limitation of PyType annotations for JAX is that they do not support the
specification ofDeviceArrayranks, shapes or dtypes. Chex includes a number
of functions that allow flexible and concise specification of these properties.E.g. suppose you want to ensure that all tensorst1,t2,t3have the same
shape, and that tensorst4,t5have rank2and (3or4), respectively.chex.assert_equal_shape([t1,t2,t3])chex.assert_rank([t4,t5],[2,{3,4}])More examples:fromcheximportassert_shape,assert_rank,...assert_shape(x,(2,3))# x has shape (2, 3)assert_shape([x,y],[(),(2,3)])# x is scalar and y has shape (2, 3)assert_rank(x,0)# x is scalarassert_rank([x,y],[0,2])# x is scalar and y is a rank-2 arrayassert_rank([x,y],{0,2})# x and y are scalar OR rank-2 arraysassert_type(x,int)# x has type `int` (x can be an array)assert_type([x,y],[int,float])# x has type `int` and y has type `float`assert_equal_shape([x,y,z])# x, y, and z have equal shapesassert_trees_all_close(tree_x,tree_y)# values and structure of trees matchassert_tree_all_finite(tree_x)# all tree_x leaves are finiteassert_devices_available(2,'gpu')# 2 GPUs availableassert_tpu_available()# at least 1 TPU availableassert_numerical_grads(f,(x,y),j)# f^{(j)}(x, y) matches numerical gradsSeeasserts.pydocumentationto
find all supported assertions.If you cannot find a specific assertion, please consider making a pull request
or openning an issue onthe bug tracker.Optional ArgumentsAll chex assertions support the following optional kwargs for manipulating the
emitted exception messages:custom_message: A string to include into the emitted exception messages.include_default_message: Whether to include the default Chex message into
the emitted exception messages.exception_type: An exception type to use.AssertionErrorby default.For example, the following code:dataset=load_dataset()params=init_params()foriinrange(num_steps):params=update_params(params,dataset.sample())chex.assert_tree_all_finite(params,custom_message=f'Failed at iteration{i}.',exception_type=ValueError)will raise aValueErrorthat includes a step number whenparamsget polluted
withNaNsorNones.Static and Value (akaRuntime) AssertionsChex divides all assertions into 2 classes:staticandvalueassertions.staticassertions use anything except concrete values of tensors.
Examples:assert_shape,assert_trees_all_equal_dtypes,assert_max_traces.valueassertions require access to tensor values, which are not
available during JAX tracing (seeHowJAX primitives work),
thus such assertion need special treatment in ajittedcode.To enable value assertions in a jitted function, it can be decorated withchex.chexify()wrapper. Example:@chex.chexify@jax.jitdeflogp1_abs_safe(x:chex.Array)->chex.Array:chex.assert_tree_all_finite(x)returnjnp.log(jnp.abs(x)+1)logp1_abs_safe(jnp.ones(2))# OKlogp1_abs_safe(jnp.array([jnp.nan,3]))# FAILS (in async mode)# The error will be raised either at the next line OR at the next# `logp1_abs_safe` call. See the docs for more detain on async mode.logp1_abs_safe.wait_checks()# Wait for the (async) computation to complete.Seethis docstringfor more detail onchex.chexify().JAX Tracing AssertionsJAX re-traces JIT'ted function every time the structure of passed arguments
changes. Often this behavior is inadvertent and leads to a significant
performance drop which is hard to debug.@chex.assert_max_tracesdecorator asserts that the function is not re-traced more thanntimes during
program execution.Global trace counter can be cleared by callingchex.clear_trace_counter(). This function be used to isolate unittests relying
on@chex.assert_max_traces.Examples:@jax.jit@chex.assert_max_traces(n=1)deffn_sum_jitted(x,y):returnx+yfn_sum_jitted(jnp.zeros(3),jnp.zeros(3))# tracing for the 1st time - OKfn_sum_jitted(jnp.zeros([6,7]),jnp.zeros([6,7]))# AssertionError!Can be used withjax.pmap()as well:deffn_sub(x,y):returnx-yfn_sub_pmapped=jax.pmap(chex.assert_max_traces(fn_sub,n=10))SeeHowJAX primitives worksection for more information about tracing.Test variants (variants.py)JAX relies extensively on code transformation and compilation, meaning that it
can be hard to ensure that code is properly tested. For instance, just testing a
python function using JAX code will not cover the actual code path that is
executed when jitted, and that path will also differ whether the code is jitted
for CPU, GPU, or TPU. This has been a source of obscure and hard to catch bugs
where XLA changes would lead to undesirable behaviours that however only
manifest in one specific code transformation.Variants make it easy to ensure that unit tests cover different ‘variations’ of
a function, by providing a simple decorator that can be used to repeat any test
under all (or a subset) of the relevant code transformations.E.g. suppose you want to test the output of a functionfnwith or without jit.
You can usechex.variantsto run the test with both the jitted and non-jitted
version of the function by simply decorating a test method with@chex.variants, and then usingself.variant(fn)in place offnin the body
of the test.deffn(x,y):returnx+y...classExampleTest(chex.TestCase):@chex.variants(with_jit=True,without_jit=True)deftest(self):var_fn=self.variant(fn)self.assertEqual(fn(1,2),3)self.assertEqual(var_fn(1,2),fn(1,2))If you define the function in the test method, you may also useself.variantas a decorator in the function definition. For example:classExampleTest(chex.TestCase):@chex.variants(with_jit=True,without_jit=True)deftest(self):@self.variantdefvar_fn(x,y):returnx+yself.assertEqual(var_fn(1,2),3)Example of parameterized test:fromabsl.testingimportparameterized# Could also be:#  `class ExampleParameterizedTest(chex.TestCase, parameterized.TestCase):`#  `class ExampleParameterizedTest(chex.TestCase):`classExampleParameterizedTest(parameterized.TestCase):@chex.variants(with_jit=True,without_jit=True)@parameterized.named_parameters(('case_positive',1,2,3),('case_negative',-1,-2,-3),)deftest(self,arg_1,arg_2,expected):@self.variantdefvar_fn(x,y):returnx+yself.assertEqual(var_fn(arg_1,arg_2),expected)Chex currently supports the following variants:with_jit-- appliesjax.jit()transformation to the function.without_jit-- uses the function as is, i.e. identity transformation.with_device-- places all arguments (except specified inignore_argnumsargument) into device memory before applying the function.without_device-- places all arguments in RAM before applying the function.with_pmap-- appliesjax.pmap()transformation to the function (see notes below).See documentation invariants.pyfor more details on the supported variants.
More examples can be found invariants_test.py.Variants notesTest classes that use@chex.variantsmust inherit fromchex.TestCase(or any other base class that unrolls tests generators
withinTestCase, e.g.absl.testing.parameterized.TestCase).[jax.vmap]All variants can be applied to a vmapped function;
please see an example invariants_test.py(test_vmapped_fn_named_paramsandtest_pmap_vmapped_fn).[@chex.all_variants]You can get all supported variants
by using the decorator@chex.all_variants.[with_pmapvariant]jax.pmap(fn)(doc) performs
parallel map offnonto multiple devices. Since most tests run in a
single-device environment (i.e. having access to a single CPU or GPU), in which
casejax.pmapis a functional equivalent tojax.jit,with_pmapvariant is
skipped by default (although it works fine with a single device). Below we
describe  a way to properly testfnif it is supposed to be used in
multi-device environments (TPUs or multiple CPUs/GPUs). To disable skippingwith_pmapvariants in case of a single device, add--chex_skip_pmap_variant_if_single_device=falseto your test command.Fakes (fake.py)Debugging in JAX is made more difficult by code transformations such asjitandpmap, which introduce optimizations that make code hard to inspect and
trace. It can also be difficult to disable those transformations during
debugging as they can be called at several places in the underlying
code. Chex provides tools to globally replacejax.jitwith a no-op
transformation andjax.pmapwith a (non-parallel)jax.vmap, in order to more
easily debug code in a single-device context.For example, you can use Chex to fakepmapand have it replaced with avmap.
This can be achieved by wrapping your code with a context manager:withchex.fake_pmap():@jax.pmapdeffn(inputs):...# Function will be vmapped over inputsfn(inputs)The same functionality can also be invoked withstartandstop:fake_pmap=chex.fake_pmap()fake_pmap.start()...yourjaxcode...fake_pmap.stop()In addition, you can fake a real multi-device test environment with a
multi-threaded CPU. See sectionFaking multi-device test environmentsfor
more details.See documentation infake.pyand examples infake_test.pyfor more details.Faking multi-device test environmentsIn situations where you do not have easy access to multiple devices, you can
still test parallel computation using single-device multi-threading.In particular, one can force XLA to use a single CPU's threads as separate
devices, i.e. to fake a real multi-device environment with a multi-threaded one.
These two options are theoretically equivalent from XLA perspective because they
expose the same interface and use identical abstractions.Chex has a flagchex_n_cpu_devicesthat specifies a number of CPU threads to
use as XLA devices.To set up a multi-threaded XLA environment forabsltests, definesetUpModulefunction in your test module:defsetUpModule():chex.set_n_cpu_devices()Now you can launch your test withpython test.py --chex_n_cpu_devices=Nto run
it in multi-device regime. Note thatalltests within a module will have an
access toNdevices.More examples can be found invariants_test.py,fake_test.pyandfake_set_n_cpu_devices_test.py.Using named dimension sizes.Chex comes with a small utility that allows you to package a collection of
dimension sizes into a single object. The basic idea is:dims=chex.Dimensions(B=batch_size,T=sequence_len,E=embedding_dim)...chex.assert_shape(arr,dims['BTE'])String lookups are translated integer tuples. For instance, let's saybatch_size == 3,sequence_len = 5andembedding_dim = 7, thendims['BTE']==(3,5,7)dims['B']==(3,)dims['TTBEE']==(5,5,3,7,7)...You can also assign dimension sizes dynamically as follows:dims['XY']=some_matrix.shapedims.Z=13For more examples, seechex.Dimensionsdocumentation.Citing ChexThis repository is part of theDeepMind JAX Ecosystem, to cite Chex please use
theDeepMind JAX Ecosystem citation.