Chex

Chex 是一个用于帮助编写可靠 JAX 代码的工具库。

它包含以下功能：

• 对代码进行插桩（例如断言、警告）
• 调试（例如在上下文管理器中将 pmap 转换为 vmap）
• 在多种变体（例如已 jitted 和未 jitted）下测试 JAX 代码

安装

你可以通过 PyPI 安装最新发布的 Chex 版本：

pip install chex

或者从 GitHub 安装最新的开发版本：

pip install git+https://github.com/deepmind/chex.git

模块概览

Dataclass (dataclass.py)

数据类是 Python 3.7 引入的一种流行构造，允许以最少的样板代码轻松定义类型化的数据结构。然而，它们默认情况下并不兼容 JAX 和 dm-tree。

在 Chex 中，我们提供了一个对 JAX 友好的数据类实现，复用了 Python 的 dataclasses。

Chex 的 dataclass 实现会将数据类注册为内部的 PyTree 节点，以确保与 JAX 数据结构的兼容性。

此外，我们还提供了一个类包装器，将数据类暴露为 collections.Mapping 的子类，从而可以在 dm-tree 方法中像处理普通 Python 字典一样对其进行（展平或反展平）操作。更多详细信息请参阅 @mappable_dataclass 的文档字符串。

示例：

@chex.dataclass
class Parameters:
  x: chex.ArrayDevice
  y: chex.ArrayDevice

parameters = Parameters(
    x=jnp.ones((2, 2)),
    y=jnp.ones((1, 2)),
)

# 数据类可以作为 JAX 的 pytree 处理
jax.tree_util.tree_map(lambda x: 2.0 * x, parameters)

# 也可以作为 dm-tree 中的映射处理
tree.flatten(parameters)

注意：与标准的 Python 3.7 数据类不同，Chex 数据类不能使用位置参数进行实例化。它们仅支持与 Python 字典构造函数格式相同的命名参数。如果需要，可以通过 from_tuple 和 to_tuple 方法将数据类转换为元组。

parameters = Parameters(
    jnp.ones((2, 2)),
    jnp.ones((1, 2)),
)
# ValueError: 可映射的数据类构造函数不支持位置参数。

断言（asserts.py）

PyType 注解在 JAX 中的一个局限性在于，它们无法指定 DeviceArray 的秩、形状或数据类型。Chex 提供了一系列函数，可以灵活且简洁地指定这些属性。

例如，假设你希望确保张量 t1、t2、t3 具有相同的形状，并且张量 t4 和 t5 分别具有 2 阶和（3 或 4）阶的形状：

chex.assert_equal_shape([t1, t2, t3])
chex.assert_rank([t4, t5], [2, {3, 4}])

更多示例：

from chex import assert_shape, assert_rank, ...

assert_shape(x, (2, 3))                # x 的形状为 (2, 3)
assert_shape([x, y], [(), (2,3)])      # x 是标量，y 的形状为 (2, 3)

assert_rank(x, 0)                      # x 是标量
assert_rank([x, y], [0, 2])            # x 是标量，y 是 2 阶数组
assert_rank([x, y], {0, 2})            # x 和 y 要么是标量，要么是 2 阶数组

assert_type(x, int)                    # x 的类型为 int（x 可以是数组）
assert_type([x, y], [int, float])      # x 的类型为 int，y 的类型为 float

assert_equal_shape([x, y, z])          # x、y 和 z 的形状相同

assert_trees_all_close(tree_x, tree_y) # 树结构及对应值完全一致
assert_tree_all_finite(tree_x)         # tree_x 的所有叶节点均为有限值

assert_devices_available(2, 'gpu')     # 有 2 张 GPU 可用
assert_tpu_available()                 # 至少有 1 张 TPU 可用

assert_numerical_grads(f, (x, y), j)   # f 的第 j 阶导数与数值梯度一致

请参阅 asserts.py 文档，以了解所有支持的断言。

如果您找不到所需的特定断言，请考虑提交拉取请求或在 问题跟踪器 上提出问题。

可选参数

所有 Chex 断言都支持以下可选关键字参数，用于自定义抛出的异常消息：

• custom_message: 用于包含在异常消息中的自定义字符串。
• include_default_message: 是否将默认的 Chex 消息也包含在异常消息中。
• exception_type: 使用的异常类型，默认为 AssertionError。

例如，以下代码：

dataset = load_dataset()
params = init_params()
for i in range(num_steps):
  params = update_params(params, dataset.sample())
  chex.assert_tree_all_finite(params,
                              custom_message=f'Failed at iteration {i}.',
                              exception_type=ValueError)

将在 params 被 NaN 或 None 污染时抛出一个包含迭代次数的 ValueError。

静态与值（运行时）断言

Chex 将所有断言分为两类：静态断言和值断言。

1. 静态断言不依赖于张量的具体数值。例如：assert_shape、assert_trees_all_equal_dtypes、assert_max_traces 等。
2. 值断言需要访问张量的数值，而这些数值在 JAX 追踪期间是不可用的（参见 JAX 原语的工作原理）， 因此这类断言在 jitted 代码中需要特殊处理。

为了在 jitted 函数中启用值断言，可以使用 chex.chexify() 包装器对其进行装饰。示例如下：

  @chex.chexify
  @jax.jit
  def logp1_abs_safe(x: chex.Array) -> chex.Array:
    chex.assert_tree_all_finite(x)
    return jnp.log(jnp.abs(x) + 1)

  logp1_abs_safe(jnp.ones(2))  # 正常
  logp1_abs_safe(jnp.array([jnp.nan, 3]))  # 失败（异步模式下）

  # 错误可能会在下一行或下次调用 `logp1_abs_safe` 时抛出。有关异步模式的更多信息，请参阅文档。
  logp1_abs_safe.wait_checks()  # 等待（异步）计算完成。

有关 chex.chexify() 的更多详细信息，请参阅 this 文档字符串。

JAX 追踪断言

JAX 会在每次传入参数的结构发生变化时重新追踪 JIT 编译的函数。这种行为通常是无意的，会导致性能显著下降，且难以调试。@chex.assert_max_traces 装饰器会断言函数在程序执行过程中不会被重新追踪超过 n 次。

可以通过调用 chex.clear_trace_counter() 来重置全局追踪计数器。此函数可用于隔离依赖于 @chex.assert_max_traces 的单元测试。

示例：

  @jax.jit
  @chex.assert_max_traces(n=1)
  def fn_sum_jitted(x, y):
    return x + y

  fn_sum_jitted(jnp.zeros(3), jnp.zeros(3))  # 第一次追踪 - 正常
  fn_sum_jitted(jnp.zeros([6, 7]), jnp.zeros([6, 7]))  # 抛出 AssertionError!

也可以与 jax.pmap() 一起使用：

  def fn_sub(x, y):
    return x - y

  fn_sub_pmapped = jax.pmap(chex.assert_max_traces(fn_sub, n=10))

有关追踪的更多信息，请参阅 如何工作 JAX 原语 一节。

警告（warnigns.py）

除了硬性断言之外，Chex 还提供了一些实用工具来添加常见的警告，例如特定类型的弃用警告。

测试变体（variants.py）

JAX 大量依赖于代码变换和编译，这意味着很难确保代码得到充分测试。例如，仅使用 JAX 代码测试一个 Python 函数，并不能覆盖在使用 jit 编译时实际执行的代码路径；而这条路径还会因代码是为 CPU、GPU 还是 TPU 编译而有所不同。这曾导致一些难以发现的晦涩 bug：XLA 的某些更改可能会引发不良行为，但这些行为只会在特定的代码变换中才会显现。

变体机制通过提供一个简单的装饰器，可以轻松地确保单元测试覆盖函数的不同“变体”，即在所有（或部分）相关的代码变换下重复运行测试。

例如，假设你想测试函数 fn 在启用和禁用 jit 编译时的输出。你可以使用 chex.variants 来分别以已 jit 和未 jit 的版本运行测试，只需在测试方法上添加 @chex.variants 装饰器，然后在测试体中用 self.variant(fn) 替代 fn 即可。

def fn(x, y):
  return x + y
...

class ExampleTest(chex.TestCase):

  @chex.variants(with_jit=True, without_jit=True)
  def test(self):
    var_fn = self.variant(fn)
    self.assertEqual(fn(1, 2), 3)
    self.assertEqual(var_fn(1, 2), fn(1, 2))

如果在测试方法中定义函数，也可以将 self.variant 作为装饰器直接用于函数定义。例如：

class ExampleTest(chex.TestCase):

  @chex.variants(with_jit=True, without_jit=True)
  def test(self):
    @self.variant
    def var_fn(x, y):
       return x + y

    self.assertEqual(var_fn(1, 2), 3)

参数化测试示例：

from absl.testing import parameterized

# 也可以这样写：
#  `class ExampleParameterizedTest(chex.TestCase, parameterized.TestCase):`
#  `class ExampleParameterizedTest(chex.TestCase):`
class ExampleParameterizedTest(parameterized.TestCase):

  @chex.variants(with_jit=True, without_jit=True)
  @parameterized.named_parameters(
      ('case_positive', 1, 2, 3),
      ('case_negative', -1, -2, -3),
  )
  def test(self, arg_1, arg_2, expected):
    @self.variant
    def var_fn(x, y):
       return x + y

    self.assertEqual(var_fn(arg_1, arg_2), expected)

Chex 目前支持以下变体：

• with_jit — 对函数应用 jax.jit() 变换。
• without_jit — 使用原样函数，即恒等变换。
• with_device — 在应用函数之前，将所有参数（除 ignore_argnums 参数中指定的外）放入设备内存。
• without_device — 在应用函数之前，将所有参数放入 RAM。
• with_pmap — 对函数应用 jax.pmap() 变换（见下文说明）。

有关支持的变体的更多详细信息，请参阅 variants.py 中的文档。更多示例可在 variants_test.py 中找到。

变体说明

• 使用 @chex.variants 的测试类必须继承自 chex.TestCase（或任何其他在 TestCase 内展开测试生成器的基类，如 absl.testing.parameterized.TestCase）。

• [jax.vmap] 所有变体均可应用于经过 vmap 变换的函数；请参阅 variants_test.py 中的示例（test_vmapped_fn_named_params 和 test_pmap_vmapped_fn）。

• [@chex.all_variants] 您可以使用 @chex.all_variants 装饰器获取所有支持的变体。

• [with_pmap 变体] jax.pmap(fn)（文档）会将 fn 并行映射到多个设备上。由于大多数测试都在单设备环境中运行（即仅访问单个 CPU 或 GPU），在这种情况下 jax.pmap 与 jax.jit 功能上等价，因此默认会跳过 with_pmap 变体（尽管它在单设备上也能正常工作）。下面我们将介绍如何正确测试那些本应在多设备环境（TPU 或多个 CPU/GPU）中使用的函数。若要在单设备环境下不跳过 with_pmap 变体，可在测试命令中添加 --chex_skip_pmap_variant_if_single_device=false。

模拟（fake.py）

JAX 中的调试因诸如 jit 和 pmap 等代码变换而变得更加困难，这些优化使得代码难以检查和跟踪。此外，在调试过程中也很难禁用这些变换，因为它们可能在底层代码的多个位置被调用。Chex 提供了工具，可以全局将 jax.jit 替换为无操作变换，将 jax.pmap 替换为（非并行的）jax.vmap，从而更方便地在单设备环境下进行调试。

例如，您可以使用 Chex 模拟 pmap，使其替换为 vmap。这可以通过将您的代码包裹在一个上下文管理器中来实现：

with chex.fake_pmap():
  @jax.pmap
  def fn(inputs):
    ...

  # 函数将对 inputs 进行 vmap 变换
  fn(inputs)

同样的功能也可以通过 start 和 stop 方法来调用：

fake_pmap = chex.fake_pmap()
fake_pmap.start()
... 您的 jax 代码 ...
fake_pmap.stop()

此外，您还可以使用多线程 CPU 模拟真实的多设备测试环境。有关更多详情，请参阅“模拟多设备测试环境”部分。

有关更多详细信息，请参阅 fake.py 中的文档以及 fake_test.py 中的示例。

模拟多设备测试环境

在无法轻松访问多个设备的情况下，您仍然可以通过单设备多线程来测试并行计算。

具体来说，可以强制 XLA 将单个 CPU 的线程视为独立的设备，即用多线程环境模拟真实的多设备环境。从 XLA 的角度来看，这两种方式在理论上是等价的，因为它们暴露了相同的接口，并使用了相同的抽象层。

Chex 提供了一个标志 chex_n_cpu_devices，用于指定用作 XLA 设备的 CPU 线程数量。

要为 absl 测试设置一个多线程的 XLA 环境，请在您的测试模块中定义 setUpModule 函数：

def setUpModule():
  chex.set_n_cpu_devices()

现在，您可以使用 python test.py --chex_n_cpu_devices=N 来启动测试，以多设备模式运行。请注意，模块中的所有测试都将能够访问 N 个设备。

更多示例可以在以下文件中找到：variants_test.py、fake_test.py 和 fake_set_n_cpu_devices_test.py。

检测错误后端的使用（restrict_backends.py）

有时，某些 JAX 代码需要在存在加速器但已被预留用于其他用途的环境中运行。通常会使用 jax.jit(..., backend='cpu') 来避免代码使用加速器，但很难手动检查整个子系统是否都严格遵守了这一规则。此时，可以使用 restrict_backends() 来检测任何遗漏的情况，并通过抛出异常来报告。

restrict_backends() 是一个上下文管理器，它会阻止 JAX 在指定的后端上进行编译。例如：

  with chex.restrict_backends(allowed=['cpu']):
    call_some_jax_subsystem_that_must_only_run_on_cpu()

  with chex.restrict_backends(forbidden=['tpu']):
    call_some_other_subsystem_that_must_not_touch_any_tpu()

使用命名维度大小

Chex 提供了一个小型工具，允许您将一组维度大小打包成一个对象。其基本思想如下：

dims = chex.Dimensions(B=batch_size, T=sequence_len, E=embedding_dim)
...
chex.assert_shape(arr, dims['BTE'])

字符串查找会被转换为整数元组。例如，假设 batch_size == 3、sequence_len = 5 和 embedding_dim = 7，则：

dims['BTE'] == (3, 5, 7)
dims['B'] == (3,)
dims['TTBEE'] == (5, 5, 3, 7, 7)
...

您还可以动态地分配维度大小，如下所示：

dims['XY'] = some_matrix.shape
dims.Z = 13

更多示例请参阅 chex.Dimensions 文档。

引用 Chex

本仓库是 [DeepMind JAX 生态系统] 的一部分。如需引用 Chex，请使用以下参考文献：

@software{deepmind2020jax,
  title = {The {D}eep{M}ind {JAX} {E}cosystem},
  author = {DeepMind and Babuschkin, Igor and Baumli, Kate and Bell, Alison and Bhupatiraju, Surya and Bruce, Jake and Buchlovsky, Peter and Budden, David and Cai, Trevor and Clark, Aidan and Danihelka, Ivo and Dedieu, Antoine and Fantacci, Claudio and Godwin, Jonathan and Jones, Chris and Hemsley, Ross and Hennigan, Tom and Hessel, Matteo and Hou, Shaobo and Kapturowski, Steven and Keck, Thomas and Kemaev, Iurii and King, Michael and Kunesch, Markus and Martens, Lena and Merzic, Hamza and Mikulik, Vladimir and Norman, Tamara and Papamakarios, George and Quan, John and Ring, Roman and Ruiz, Francisco and Sanchez, Alvaro and Sartran, Laurent and Schneider, Rosalia and Sezener, Eren and Spencer, Stephen and Srinivasan, Srivatsan and Stanojevi\'{c}, Milo\v{s} and Stokowiec, Wojciech and Wang, Luyu and Zhou, Guangyao and Viola, Fabio},
  url = {http://github.com/deepmind},
  year = {2020},
}

Chex 快速上手指南

Chex 是一个专为 JAX 设计的实用工具库，旨在帮助开发者编写更可靠、更易调试的 JAX 代码。它提供了类型检查、断言验证、测试变体（Variants）以及兼容 JAX PyTree 的数据类等功能。

环境准备

在安装 Chex 之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS, 或 Windows (WSL 推荐)
• Python 版本: Python 3.8+
• 核心依赖:
  • JAX (必须已安装)
  • dm-tree (通常随 JAX 生态自动安装，若缺失需手动安装)
  • Abseil (用于测试工具)

提示：如果您尚未安装 JAX，建议先根据 JAX 官方文档 完成 JAX 及其加速后端（CPU/GPU/TPU）的安装。

安装步骤

您可以选择从 PyPI 安装稳定版，或从 GitHub 安装最新开发版。

方式一：安装稳定版（推荐）

使用 pip 直接安装：

pip install chex

国内加速方案： 如果您在中国大陆地区，建议使用国内镜像源以加快下载速度：

pip install chex -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：安装开发版

如需体验最新功能或修复：

pip install git+https://github.com/deepmind/chex.git

国内加速：

pip install git+https://github.com/deepmind/chex.git -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

Chex 的核心功能主要包括：JAX 友好的数据类、运行时断言以及多变体测试。以下是三个最常用的场景示例。

1. 使用 JAX 兼容的 Dataclass

Chex 提供的 @chex.dataclass 装饰器可以将 Python 数据类注册为 JAX PyTree 节点，使其能够直接被 jax.tree_util 处理，同时也兼容 dm-tree。

import jax.numpy as jnp
import jax
import chex

@chex.dataclass
class Parameters:
  x: chex.ArrayDevice
  y: chex.ArrayDevice

# 初始化参数 (注意：必须使用关键字参数，不支持位置参数)
params = Parameters(
    x=jnp.ones((2, 2)),
    y=jnp.ones((1, 2)),
)

# 作为 JAX PyTree 进行操作
doubled_params = jax.tree_util.tree_map(lambda x: 2.0 * x, params)

# 作为 Mapping (字典) 进行扁平化处理
import tree
flat_list, treedef = tree.flatten(params)

2. 使用断言验证张量属性

在调试 JAX 代码时，标准的类型注解无法检查 Shape、Rank 或 dtype。Chex 提供了一系列断言函数来在运行时验证这些属性。

import chex
import jax.numpy as jnp

x = jnp.zeros((2, 3))
y = jnp.zeros((2, 3))
z = jnp.array([1.0, 2.0])

# 验证形状是否一致
chex.assert_equal_shape([x, y])

# 验证特定形状
chex.assert_shape(x, (2, 3))

# 验证秩 (Rank)
chex.assert_rank(z, 1) 

# 验证数值是否有限 (无 NaN 或 Inf)
chex.assert_tree_all_finite({"weights": x, "bias": z})

# 在 JIT 编译函数中使用断言 (需配合 chexify)
@chex.chexify
@jax.jit
def safe_compute(data):
    chex.assert_tree_all_finite(data)
    return jnp.sum(data)

safe_compute(x)

3. 使用 Variants 进行多模式测试

JAX 代码在 jit 编译前后或在不同设备上的行为可能不同。Chex 的 variants 装饰器可以自动让同一个测试用例在多种配置下运行（例如：有/无 JIT，有/无设备内存）。

import chex
import jax.numpy as jnp

def my_function(x, y):
    return x + y

class MyTest(chex.TestCase):

  # 定义测试变体：同时测试 jit 编译版和非编译版
  @chex.variants(with_jit=True, without_jit=True)
  def test_addition(self):
    # self.variant 会根据当前变体自动包装函数
    tested_fn = self.variant(my_function)
    
    x = jnp.array([1.0, 2.0])
    y = jnp.array([3.0, 4.0])
    
    result = tested_fn(x, y)
    chex.assert_equal_shape([result, x])
    chex.assert_trees_all_close(result, jnp.array([4.0, 6.0]))

# 运行测试 (通常在命令行通过 pytest 或 python -m unittest 运行)
# 此测试会自动运行两次：一次纯 Python，一次 JIT 编译

通过以上基础功能，您可以显著提升 JAX 项目的代码健壮性和测试覆盖率。更多高级用法（如 assert_max_traces 防止重复追踪）请参考官方 API 文档。