[{"data":1,"prerenderedAt":-1},["ShallowReactive",2],{"similar-google-deepmind--chex":3,"tool-google-deepmind--chex":61},[4,18,28,36,45,53],{"id":5,"name":6,"github_repo":7,"description_zh":8,"stars":9,"difficulty_score":10,"last_commit_at":11,"category_tags":12,"status":17},4358,"openclaw","openclaw\u002Fopenclaw","OpenClaw 是一款专为个人打造的本地化 AI 助手，旨在让你在自己的设备上拥有完全可控的智能伙伴。它打破了传统 AI 助手局限于特定网页或应用的束缚，能够直接接入你日常使用的各类通讯渠道，包括微信、WhatsApp、Telegram、Discord、iMessage 等数十种平台。无论你在哪个聊天软件中发送消息，OpenClaw 都能即时响应，甚至支持在 macOS、iOS 和 Android 设备上进行语音交互，并提供实时的画布渲染功能供你操控。\n\n这款工具主要解决了用户对数据隐私、响应速度以及“始终在线”体验的需求。通过将 AI 部署在本地，用户无需依赖云端服务即可享受快速、私密的智能辅助，真正实现了“你的数据，你做主”。其独特的技术亮点在于强大的网关架构，将控制平面与核心助手分离，确保跨平台通信的流畅性与扩展性。\n\nOpenClaw 非常适合希望构建个性化工作流的技术爱好者、开发者，以及注重隐私保护且不愿被单一生态绑定的普通用户。只要具备基础的终端操作能力（支持 macOS、Linux 及 Windows WSL2），即可通过简单的命令行引导完成部署。如果你渴望拥有一个懂你",349277,3,"2026-04-06T06:32:30",[13,14,15,16],"Agent","开发框架","图像","数据工具","ready",{"id":19,"name":20,"github_repo":21,"description_zh":22,"stars":23,"difficulty_score":24,"last_commit_at":25,"category_tags":26,"status":17},9989,"n8n","n8n-io\u002Fn8n","n8n 是一款面向技术团队的公平代码（fair-code）工作流自动化平台，旨在让用户在享受低代码快速构建便利的同时，保留编写自定义代码的灵活性。它主要解决了传统自动化工具要么过于封闭难以扩展、要么完全依赖手写代码效率低下的痛点，帮助用户轻松连接 400 多种应用与服务，实现复杂业务流程的自动化。\n\nn8n 特别适合开发者、工程师以及具备一定技术背景的业务人员使用。其核心亮点在于“按需编码”：既可以通过直观的可视化界面拖拽节点搭建流程，也能随时插入 JavaScript 或 Python 代码、调用 npm 包来处理复杂逻辑。此外，n8n 原生集成了基于 LangChain 的 AI 能力，支持用户利用自有数据和模型构建智能体工作流。在部署方面，n8n 提供极高的自由度，支持完全自托管以保障数据隐私和控制权，也提供云端服务选项。凭借活跃的社区生态和数百个现成模板，n8n 让构建强大且可控的自动化系统变得简单高效。",184740,2,"2026-04-19T23:22:26",[16,14,13,15,27],"插件",{"id":29,"name":30,"github_repo":31,"description_zh":32,"stars":33,"difficulty_score":10,"last_commit_at":34,"category_tags":35,"status":17},3808,"stable-diffusion-webui","AUTOMATIC1111\u002Fstable-diffusion-webui","stable-diffusion-webui 是一个基于 Gradio 构建的网页版操作界面，旨在让用户能够轻松地在本地运行和使用强大的 Stable Diffusion 图像生成模型。它解决了原始模型依赖命令行、操作门槛高且功能分散的痛点，将复杂的 AI 绘图流程整合进一个直观易用的图形化平台。\n\n无论是希望快速上手的普通创作者、需要精细控制画面细节的设计师，还是想要深入探索模型潜力的开发者与研究人员，都能从中获益。其核心亮点在于极高的功能丰富度：不仅支持文生图、图生图、局部重绘（Inpainting）和外绘（Outpainting）等基础模式，还独创了注意力机制调整、提示词矩阵、负向提示词以及“高清修复”等高级功能。此外，它内置了 GFPGAN 和 CodeFormer 等人脸修复工具，支持多种神经网络放大算法，并允许用户通过插件系统无限扩展能力。即使是显存有限的设备，stable-diffusion-webui 也提供了相应的优化选项，让高质量的 AI 艺术创作变得触手可及。",162132,"2026-04-05T11:01:52",[14,15,13],{"id":37,"name":38,"github_repo":39,"description_zh":40,"stars":41,"difficulty_score":24,"last_commit_at":42,"category_tags":43,"status":17},1381,"everything-claude-code","affaan-m\u002Feverything-claude-code","everything-claude-code 是一套专为 AI 编程助手（如 Claude Code、Codex、Cursor 等）打造的高性能优化系统。它不仅仅是一组配置文件，而是一个经过长期实战打磨的完整框架，旨在解决 AI 代理在实际开发中面临的效率低下、记忆丢失、安全隐患及缺乏持续学习能力等核心痛点。\n\n通过引入技能模块化、直觉增强、记忆持久化机制以及内置的安全扫描功能，everything-claude-code 能显著提升 AI 在复杂任务中的表现，帮助开发者构建更稳定、更智能的生产级 AI 代理。其独特的“研究优先”开发理念和针对 Token 消耗的优化策略，使得模型响应更快、成本更低，同时有效防御潜在的攻击向量。\n\n这套工具特别适合软件开发者、AI 研究人员以及希望深度定制 AI 工作流的技术团队使用。无论您是在构建大型代码库，还是需要 AI 协助进行安全审计与自动化测试，everything-claude-code 都能提供强大的底层支持。作为一个曾荣获 Anthropic 黑客大奖的开源项目，它融合了多语言支持与丰富的实战钩子（hooks），让 AI 真正成长为懂上",161147,"2026-04-19T23:31:47",[14,13,44],"语言模型",{"id":46,"name":47,"github_repo":48,"description_zh":49,"stars":50,"difficulty_score":24,"last_commit_at":51,"category_tags":52,"status":17},2271,"ComfyUI","Comfy-Org\u002FComfyUI","ComfyUI 是一款功能强大且高度模块化的视觉 AI 引擎，专为设计和执行复杂的 Stable Diffusion 图像生成流程而打造。它摒弃了传统的代码编写模式，采用直观的节点式流程图界面，让用户通过连接不同的功能模块即可构建个性化的生成管线。\n\n这一设计巧妙解决了高级 AI 绘图工作流配置复杂、灵活性不足的痛点。用户无需具备编程背景，也能自由组合模型、调整参数并实时预览效果，轻松实现从基础文生图到多步骤高清修复等各类复杂任务。ComfyUI 拥有极佳的兼容性，不仅支持 Windows、macOS 和 Linux 全平台，还广泛适配 NVIDIA、AMD、Intel 及苹果 Silicon 等多种硬件架构，并率先支持 SDXL、Flux、SD3 等前沿模型。\n\n无论是希望深入探索算法潜力的研究人员和开发者，还是追求极致创作自由度的设计师与资深 AI 绘画爱好者，ComfyUI 都能提供强大的支持。其独特的模块化架构允许社区不断扩展新功能，使其成为当前最灵活、生态最丰富的开源扩散模型工具之一，帮助用户将创意高效转化为现实。",109154,"2026-04-18T11:18:24",[14,15,13],{"id":54,"name":55,"github_repo":56,"description_zh":57,"stars":58,"difficulty_score":24,"last_commit_at":59,"category_tags":60,"status":17},6121,"gemini-cli","google-gemini\u002Fgemini-cli","gemini-cli 是一款由谷歌推出的开源 AI 命令行工具，它将强大的 Gemini 大模型能力直接集成到用户的终端环境中。对于习惯在命令行工作的开发者而言，它提供了一条从输入提示词到获取模型响应的最短路径，无需切换窗口即可享受智能辅助。\n\n这款工具主要解决了开发过程中频繁上下文切换的痛点，让用户能在熟悉的终端界面内直接完成代码理解、生成、调试以及自动化运维任务。无论是查询大型代码库、根据草图生成应用，还是执行复杂的 Git 操作，gemini-cli 都能通过自然语言指令高效处理。\n\n它特别适合广大软件工程师、DevOps 人员及技术研究人员使用。其核心亮点包括支持高达 100 万 token 的超长上下文窗口，具备出色的逻辑推理能力；内置 Google 搜索、文件操作及 Shell 命令执行等实用工具；更独特的是，它支持 MCP（模型上下文协议），允许用户灵活扩展自定义集成，连接如图像生成等外部能力。此外，个人谷歌账号即可享受免费的额度支持，且项目基于 Apache 2.0 协议完全开源，是提升终端工作效率的理想助手。",100752,"2026-04-10T01:20:03",[27,13,15,14],{"id":62,"github_repo":63,"name":64,"description_en":65,"description_zh":66,"ai_summary_zh":67,"readme_en":68,"readme_zh":69,"quickstart_zh":70,"use_case_zh":71,"hero_image_url":72,"owner_login":73,"owner_name":74,"owner_avatar_url":75,"owner_bio":76,"owner_company":65,"owner_location":65,"owner_email":65,"owner_twitter":65,"owner_website":77,"owner_url":78,"languages":79,"stars":88,"forks":89,"last_commit_at":90,"license":91,"difficulty_score":92,"env_os":93,"env_gpu":94,"env_ram":93,"env_deps":95,"category_tags":102,"github_topics":65,"view_count":24,"oss_zip_url":65,"oss_zip_packed_at":65,"status":17,"created_at":103,"updated_at":104,"faqs":105,"releases":136},9925,"google-deepmind\u002Fchex","chex",null,"Chex 是 DeepMind 专为 JAX 框架打造的一套实用工具库，旨在帮助开发者编写更可靠、健壮的代码。在利用 JAX 进行高性能数值计算时，开发者常面临类型检查不足、调试困难以及测试覆盖不全等挑战，Chex 正是为了解决这些痛点而生。\n\n它主要提供三大核心能力：首先，通过丰富的断言函数，让开发者能轻松验证张量的形状、秩（rank）和数据类型，弥补了 JAX 原生类型注解在维度细节上的缺失；其次，内置便捷的调试工具，例如能在上下文管理器中将并行映射（pmap）临时转换为向量映射（vmap），极大简化了调试流程；最后，支持跨多种变体（如即时编译与非编译模式）自动化测试代码，确保逻辑一致性。\n\nChex 的一大技术亮点是其对 Python 数据类（dataclass）的增强实现。它将数据类注册为 JAX 原生的 PyTree 节点，并兼容 dm-tree 库，使得自定义数据结构既能享受 JAX 的函数式变换，又能像字典一样被灵活处理。\n\n这款工具非常适合使用 JAX 进行深度学习模型研发的研究人员、算法工程师及开源贡献者。如果你希望提升代码质量，减少因维度不匹配或数值异常导致的隐蔽错","Chex 是 DeepMind 专为 JAX 框架打造的一套实用工具库，旨在帮助开发者编写更可靠、健壮的代码。在利用 JAX 进行高性能数值计算时，开发者常面临类型检查不足、调试困难以及测试覆盖不全等挑战，Chex 正是为了解决这些痛点而生。\n\n它主要提供三大核心能力：首先，通过丰富的断言函数，让开发者能轻松验证张量的形状、秩（rank）和数据类型，弥补了 JAX 原生类型注解在维度细节上的缺失；其次，内置便捷的调试工具，例如能在上下文管理器中将并行映射（pmap）临时转换为向量映射（vmap），极大简化了调试流程；最后，支持跨多种变体（如即时编译与非编译模式）自动化测试代码，确保逻辑一致性。\n\nChex 的一大技术亮点是其对 Python 数据类（dataclass）的增强实现。它将数据类注册为 JAX 原生的 PyTree 节点，并兼容 dm-tree 库，使得自定义数据结构既能享受 JAX 的函数式变换，又能像字典一样被灵活处理。\n\n这款工具非常适合使用 JAX 进行深度学习模型研发的研究人员、算法工程师及开源贡献者。如果你希望提升代码质量，减少因维度不匹配或数值异常导致的隐蔽错误，Chex 将成为你开发流程中得力的助手。","# Chex\n\n\u003C!-- disableFinding(SNIPPET_INVALID_LANGUAGE) -->\n\n![CI status](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fworkflows\u002Fci\u002Fbadge.svg)\n![docs](https:\u002F\u002Foss.gittoolsai.com\u002Fimages\u002Fgoogle-deepmind_chex_readme_13d664e1afd7.png)\n![pypi](https:\u002F\u002Fimg.shields.io\u002Fpypi\u002Fv\u002Fchex)\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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex.git\n```\n\n## Modules Overview\n\n### Dataclass ([dataclass.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fdataclass.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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Ftree) out of the box.\n\nIn Chex we provide a JAX-friendly dataclass implementation reusing python [dataclasses](https:\u002F\u002Fdocs.python.org\u002F3\u002Flibrary\u002Fdataclasses.html#module-dataclasses).\n\nChex implementation of `dataclass` registers dataclasses as internal [_PyTree_\nnodes](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Fpytrees.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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fdataclass.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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fasserts.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:\u002F\u002Fchex.readthedocs.io\u002Fen\u002Flatest\u002Fapi.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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fissues).\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:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Fnotebooks\u002FHow_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) -> 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:\u002F\u002Fchex.readthedocs.io\u002Fen\u002Flatest\u002Fapi.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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fasserts.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:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Fnotebooks\u002FHow_JAX_primitives_work.html)\nsection for more information about tracing.\n\n### Warnings ([warnigns.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fwarnings.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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants.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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants.py) for more details on the supported variants.\nMore examples can be found in [variants_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants_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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants_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:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Fjax.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\u002FGPUs). 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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake.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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake.py) and examples in [fake_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake_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:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants_test.py), [fake_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake_test.py) and [fake_set_n_cpu_devices_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake_set_n_cpu_devices_test.py).\n\n### Detecting Use of the Wrong Back-End ([restrict_backends.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Frestrict_backends.py))\n\nSometimes, certain JAX code needs to run in an environment where an accelerator\nis present but reserved for other purposes. Typically one would use\n`jax.jit(..., backend='cpu')` to keep the code away from the accelerator, but\nit is hard to check by hand that this has been done without exception\nthroughout an entire subsystem. Then, `restrict_backends()` can be used to\ndetect any overlooked case and report it by raising an exception.\n\n`restrict_backends()` is a context manager that objects to JAX compilation for\nspecified backends. For example:\n\n```python\n  with chex.restrict_backends(allowed=['cpu']):\n    call_some_jax_subsystem_that_must_only_run_on_cpu()\n\n  with chex.restrict_backends(forbidden=['tpu']):\n    call_some_other_subsystem_that_must_not_touch_any_tpu()\n```\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:\u002F\u002Fchex.readthedocs.io\u002Fen\u002Flatest\u002Fapi.html#chex.Dimensions)\ndocumentation.\n\n## Citing Chex\n\nThis repository is part of the [DeepMind JAX Ecosystem], to cite Chex please use\nthe following citation:\n\n```bibtex\n@software{deepmind2020jax,\n  title = {The {D}eep{M}ind {JAX} {E}cosystem},\n  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},\n  url = {http:\u002F\u002Fgithub.com\u002Fdeepmind},\n  year = {2020},\n}\n```\n","# Chex\n\n\u003C!-- disableFinding(SNIPPET_INVALID_LANGUAGE) -->\n\n![CI 状态](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fworkflows\u002Fci\u002Fbadge.svg)\n![文档](https:\u002F\u002Foss.gittoolsai.com\u002Fimages\u002Fgoogle-deepmind_chex_readme_13d664e1afd7.png)\n![PyPI](https:\u002F\u002Fimg.shields.io\u002Fpypi\u002Fv\u002Fchex)\n\nChex 是一个用于帮助编写可靠 JAX 代码的工具库。\n\n它包含以下功能：\n\n* 对代码进行插桩（例如断言、警告）\n* 调试（例如在上下文管理器中将 `pmap` 转换为 `vmap`）\n* 在多种变体（例如已 jitted 和未 jitted）下测试 JAX 代码\n\n## 安装\n\n你可以通过 PyPI 安装最新发布的 Chex 版本：\n\n```sh\npip install chex\n```\n\n或者从 GitHub 安装最新的开发版本：\n\n```sh\npip install git+https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex.git\n```\n\n## 模块概览\n\n### Dataclass ([dataclass.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fdataclass.py))\n\n数据类是 Python 3.7 引入的一种流行构造，允许以最少的样板代码轻松定义类型化的数据结构。然而，它们默认情况下并不兼容 JAX 和 `dm-tree`。\n\n在 Chex 中，我们提供了一个对 JAX 友好的数据类实现，复用了 Python 的 [dataclasses](https:\u002F\u002Fdocs.python.org\u002F3\u002Flibrary\u002Fdataclasses.html#module-dataclasses)。\n\nChex 的 `dataclass` 实现会将数据类注册为内部的 [_PyTree_ 节点](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Fpytrees.html)，以确保与 JAX 数据结构的兼容性。\n\n此外，我们还提供了一个类包装器，将数据类暴露为 `collections.Mapping` 的子类，从而可以在 `dm-tree` 方法中像处理普通 Python 字典一样对其进行（展平或反展平）操作。更多详细信息请参阅 [`@mappable_dataclass`](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fdataclass.py#L27) 的文档字符串。\n\n示例：\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# 数据类可以作为 JAX 的 pytree 处理\njax.tree_util.tree_map(lambda x: 2.0 * x, parameters)\n\n# 也可以作为 dm-tree 中的映射处理\ntree.flatten(parameters)\n```\n\n**注意**：与标准的 Python 3.7 数据类不同，Chex 数据类不能使用位置参数进行实例化。它们仅支持与 Python 字典构造函数格式相同的命名参数。如果需要，可以通过 `from_tuple` 和 `to_tuple` 方法将数据类转换为元组。\n\n```python\nparameters = Parameters(\n    jnp.ones((2, 2)),\n    jnp.ones((1, 2)),\n)\n# ValueError: 可映射的数据类构造函数不支持位置参数。\n```\n\n### 断言（[asserts.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fasserts.py)）\n\nPyType 注解在 JAX 中的一个局限性在于，它们无法指定 `DeviceArray` 的秩、形状或数据类型。Chex 提供了一系列函数，可以灵活且简洁地指定这些属性。\n\n例如，假设你希望确保张量 `t1`、`t2`、`t3` 具有相同的形状，并且张量 `t4` 和 `t5` 分别具有 2 阶和（3 或 4）阶的形状：\n\n```python\nchex.assert_equal_shape([t1, t2, t3])\nchex.assert_rank([t4, t5], [2, {3, 4}])\n```\n\n更多示例：\n\n```python\nfrom chex import assert_shape, assert_rank, ...\n\nassert_shape(x, (2, 3))                # x 的形状为 (2, 3)\nassert_shape([x, y], [(), (2,3)])      # x 是标量，y 的形状为 (2, 3)\n\nassert_rank(x, 0)                      # x 是标量\nassert_rank([x, y], [0, 2])            # x 是标量，y 是 2 阶数组\nassert_rank([x, y], {0, 2})            # x 和 y 要么是标量，要么是 2 阶数组\n\nassert_type(x, int)                    # x 的类型为 int（x 可以是数组）\nassert_type([x, y], [int, float])      # x 的类型为 int，y 的类型为 float\n\nassert_equal_shape([x, y, z])          # x、y 和 z 的形状相同\n\nassert_trees_all_close(tree_x, tree_y) # 树结构及对应值完全一致\nassert_tree_all_finite(tree_x)         # tree_x 的所有叶节点均为有限值\n\nassert_devices_available(2, 'gpu')     # 有 2 张 GPU 可用\nassert_tpu_available()                 # 至少有 1 张 TPU 可用\n\nassert_numerical_grads(f, (x, y), j)   # f 的第 j 阶导数与数值梯度一致\n```\n\n请参阅 `asserts.py`\n[文档](https:\u002F\u002Fchex.readthedocs.io\u002Fen\u002Flatest\u002Fapi.html#assertions)，以了解所有支持的断言。\n\n如果您找不到所需的特定断言，请考虑提交拉取请求或在\n[问题跟踪器](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fissues) 上提出问题。\n\n#### 可选参数\n\n所有 Chex 断言都支持以下可选关键字参数，用于自定义抛出的异常消息：\n\n* `custom_message`: 用于包含在异常消息中的自定义字符串。\n* `include_default_message`: 是否将默认的 Chex 消息也包含在异常消息中。\n* `exception_type`: 使用的异常类型，默认为 `AssertionError`。\n\n例如，以下代码：\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\n将在 `params` 被 `NaN` 或 `None` 污染时抛出一个包含迭代次数的 `ValueError`。\n\n#### 静态与值（运行时）断言\n\nChex 将所有断言分为两类：***静态***断言和***值***断言。\n\n1. ***静态***断言不依赖于张量的具体数值。例如：`assert_shape`、`assert_trees_all_equal_dtypes`、`assert_max_traces` 等。\n2. ***值***断言需要访问张量的数值，而这些数值在 JAX 追踪期间是不可用的（参见\n   [JAX 原语的工作原理](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Fnotebooks\u002FHow_JAX_primitives_work.html)），\n   因此这类断言在 *jitted* 代码中需要特殊处理。\n\n为了在 jitted 函数中启用值断言，可以使用 `chex.chexify()` 包装器对其进行装饰。示例如下：\n\n```python\n  @chex.chexify\n  @jax.jit\n  def logp1_abs_safe(x: chex.Array) -> 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))  # 正常\n  logp1_abs_safe(jnp.array([jnp.nan, 3]))  # 失败（异步模式下）\n\n  # 错误可能会在下一行或下次调用 `logp1_abs_safe` 时抛出。有关异步模式的更多信息，请参阅文档。\n  logp1_abs_safe.wait_checks()  # 等待（异步）计算完成。\n```\n\n有关 `chex.chexify()` 的更多详细信息，请参阅\n[this 文档字符串](https:\u002F\u002Fchex.readthedocs.io\u002Fen\u002Flatest\u002Fapi.html#chex.chexify)。\n\n#### JAX 追踪断言\n\nJAX 会在每次传入参数的结构发生变化时重新追踪 JIT 编译的函数。这种行为通常是无意的，会导致性能显著下降，且难以调试。[@chex.assert_max_traces](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fasserts.py#L44) 装饰器会断言函数在程序执行过程中不会被重新追踪超过 `n` 次。\n\n可以通过调用 `chex.clear_trace_counter()` 来重置全局追踪计数器。此函数可用于隔离依赖于 `@chex.assert_max_traces` 的单元测试。\n\n示例：\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))  # 第一次追踪 - 正常\n  fn_sum_jitted(jnp.zeros([6, 7]), jnp.zeros([6, 7]))  # 抛出 AssertionError!\n```\n\n也可以与 `jax.pmap()` 一起使用：\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\n有关追踪的更多信息，请参阅\n[如何工作 JAX 原语](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Fnotebooks\u002FHow_JAX_primitives_work.html) 一节。\n\n### 警告（[warnigns.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fwarnings.py)）\n\n除了硬性断言之外，Chex 还提供了一些实用工具来添加常见的警告，例如特定类型的弃用警告。\n\n### 测试变体（[variants.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants.py)）\n\nJAX 大量依赖于代码变换和编译，这意味着很难确保代码得到充分测试。例如，仅使用 JAX 代码测试一个 Python 函数，并不能覆盖在使用 `jit` 编译时实际执行的代码路径；而这条路径还会因代码是为 CPU、GPU 还是 TPU 编译而有所不同。这曾导致一些难以发现的晦涩 bug：XLA 的某些更改可能会引发不良行为，但这些行为只会在特定的代码变换中才会显现。\n\n变体机制通过提供一个简单的装饰器，可以轻松地确保单元测试覆盖函数的不同“变体”，即在所有（或部分）相关的代码变换下重复运行测试。\n\n例如，假设你想测试函数 `fn` 在启用和禁用 `jit` 编译时的输出。你可以使用 `chex.variants` 来分别以已 `jit` 和未 `jit` 的版本运行测试，只需在测试方法上添加 `@chex.variants` 装饰器，然后在测试体中用 `self.variant(fn)` 替代 `fn` 即可。\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\n如果在测试方法中定义函数，也可以将 `self.variant` 作为装饰器直接用于函数定义。例如：\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\n参数化测试示例：\n\n```python\nfrom absl.testing import parameterized\n\n# 也可以这样写：\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 目前支持以下变体：\n\n* `with_jit` — 对函数应用 `jax.jit()` 变换。\n* `without_jit` — 使用原样函数，即恒等变换。\n* `with_device` — 在应用函数之前，将所有参数（除 `ignore_argnums` 参数中指定的外）放入设备内存。\n* `without_device` — 在应用函数之前，将所有参数放入 RAM。\n* `with_pmap` — 对函数应用 `jax.pmap()` 变换（见下文说明）。\n\n有关支持的变体的更多详细信息，请参阅 [variants.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants.py) 中的文档。更多示例可在 [variants_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants_test.py) 中找到。\n\n### 变体说明\n\n* 使用 `@chex.variants` 的测试类必须继承自 `chex.TestCase`（或任何其他在 `TestCase` 内展开测试生成器的基类，如 `absl.testing.parameterized.TestCase`）。\n\n* **[`jax.vmap`]** 所有变体均可应用于经过 `vmap` 变换的函数；请参阅 [variants_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants_test.py) 中的示例（`test_vmapped_fn_named_params` 和 `test_pmap_vmapped_fn`）。\n\n* **[`@chex.all_variants`]** 您可以使用 `@chex.all_variants` 装饰器获取所有支持的变体。\n\n* **[`with_pmap` 变体]** `jax.pmap(fn)`（[文档](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Fjax.html#jax.pmap)）会将 `fn` 并行映射到多个设备上。由于大多数测试都在单设备环境中运行（即仅访问单个 CPU 或 GPU），在这种情况下 `jax.pmap` 与 `jax.jit` 功能上等价，因此默认会跳过 `with_pmap` 变体（尽管它在单设备上也能正常工作）。下面我们将介绍如何正确测试那些本应在多设备环境（TPU 或多个 CPU\u002FGPU）中使用的函数。若要在单设备环境下不跳过 `with_pmap` 变体，可在测试命令中添加 `--chex_skip_pmap_variant_if_single_device=false`。\n\n### 模拟（[fake.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake.py)）\n\nJAX 中的调试因诸如 `jit` 和 `pmap` 等代码变换而变得更加困难，这些优化使得代码难以检查和跟踪。此外，在调试过程中也很难禁用这些变换，因为它们可能在底层代码的多个位置被调用。Chex 提供了工具，可以全局将 `jax.jit` 替换为无操作变换，将 `jax.pmap` 替换为（非并行的）`jax.vmap`，从而更方便地在单设备环境下进行调试。\n\n例如，您可以使用 Chex 模拟 `pmap`，使其替换为 `vmap`。这可以通过将您的代码包裹在一个上下文管理器中来实现：\n\n```python\nwith chex.fake_pmap():\n  @jax.pmap\n  def fn(inputs):\n    ...\n\n  # 函数将对 inputs 进行 vmap 变换\n  fn(inputs)\n```\n\n同样的功能也可以通过 `start` 和 `stop` 方法来调用：\n\n```python\nfake_pmap = chex.fake_pmap()\nfake_pmap.start()\n... 您的 jax 代码 ...\nfake_pmap.stop()\n```\n\n此外，您还可以使用多线程 CPU 模拟真实的多设备测试环境。有关更多详情，请参阅“模拟多设备测试环境”部分。\n\n有关更多详细信息，请参阅 [fake.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake.py) 中的文档以及 [fake_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake_test.py) 中的示例。\n\n## 模拟多设备测试环境\n\n在无法轻松访问多个设备的情况下，您仍然可以通过单设备多线程来测试并行计算。\n\n具体来说，可以强制 XLA 将单个 CPU 的线程视为独立的设备，即用多线程环境模拟真实的多设备环境。从 XLA 的角度来看，这两种方式在理论上是等价的，因为它们暴露了相同的接口，并使用了相同的抽象层。\n\nChex 提供了一个标志 `chex_n_cpu_devices`，用于指定用作 XLA 设备的 CPU 线程数量。\n\n要为 `absl` 测试设置一个多线程的 XLA 环境，请在您的测试模块中定义 `setUpModule` 函数：\n\n```python\ndef setUpModule():\n  chex.set_n_cpu_devices()\n```\n\n现在，您可以使用 `python test.py --chex_n_cpu_devices=N` 来启动测试，以多设备模式运行。请注意，**模块中的所有测试**都将能够访问 `N` 个设备。\n\n更多示例可以在以下文件中找到：[variants_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Fvariants_test.py)、[fake_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake_test.py) 和 [fake_set_n_cpu_devices_test.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Ffake_set_n_cpu_devices_test.py)。\n\n### 检测错误后端的使用（[restrict_backends.py](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fblob\u002Fmaster\u002Fchex\u002F_src\u002Frestrict_backends.py)）\n\n有时，某些 JAX 代码需要在存在加速器但已被预留用于其他用途的环境中运行。通常会使用 `jax.jit(..., backend='cpu')` 来避免代码使用加速器，但很难手动检查整个子系统是否都严格遵守了这一规则。此时，可以使用 `restrict_backends()` 来检测任何遗漏的情况，并通过抛出异常来报告。\n\n`restrict_backends()` 是一个上下文管理器，它会阻止 JAX 在指定的后端上进行编译。例如：\n\n```python\n  with chex.restrict_backends(allowed=['cpu']):\n    call_some_jax_subsystem_that_must_only_run_on_cpu()\n\n  with chex.restrict_backends(forbidden=['tpu']):\n    call_some_other_subsystem_that_must_not_touch_any_tpu()\n```\n\n### 使用命名维度大小\n\nChex 提供了一个小型工具，允许您将一组维度大小打包成一个对象。其基本思想如下：\n\n```python\ndims = chex.Dimensions(B=batch_size, T=sequence_len, E=embedding_dim)\n...\nchex.assert_shape(arr, dims['BTE'])\n```\n\n字符串查找会被转换为整数元组。例如，假设 `batch_size == 3`、`sequence_len = 5` 和 `embedding_dim = 7`，则：\n\n```python\ndims['BTE'] == (3, 5, 7)\ndims['B'] == (3,)\ndims['TTBEE'] == (5, 5, 3, 7, 7)\n...\n```\n\n您还可以动态地分配维度大小，如下所示：\n\n```python\ndims['XY'] = some_matrix.shape\ndims.Z = 13\n```\n\n更多示例请参阅 [chex.Dimensions](https:\u002F\u002Fchex.readthedocs.io\u002Fen\u002Flatest\u002Fapi.html#chex.Dimensions) 文档。\n\n## 引用 Chex\n\n本仓库是 [DeepMind JAX 生态系统] 的一部分。如需引用 Chex，请使用以下参考文献：\n\n```bibtex\n@software{deepmind2020jax,\n  title = {The {D}eep{M}ind {JAX} {E}cosystem},\n  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},\n  url = {http:\u002F\u002Fgithub.com\u002Fdeepmind},\n  year = {2020},\n}\n```","# Chex 快速上手指南\n\nChex 是一个专为 JAX 设计的实用工具库，旨在帮助开发者编写更可靠、更易调试的 JAX 代码。它提供了类型检查、断言验证、测试变体（Variants）以及兼容 JAX PyTree 的数据类等功能。\n\n## 环境准备\n\n在安装 Chex 之前，请确保您的开发环境满足以下要求：\n\n*   **操作系统**: Linux, macOS, 或 Windows (WSL 推荐)\n*   **Python 版本**: Python 3.8+\n*   **核心依赖**:\n    *   [JAX](https:\u002F\u002Fgithub.com\u002Fgoogle\u002Fjax) (必须已安装)\n    *   [dm-tree](https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Ftree) (通常随 JAX 生态自动安装，若缺失需手动安装)\n    *   [Abseil](https:\u002F\u002Fgithub.com\u002Fabseil\u002Fabseil-py) (用于测试工具)\n\n> **提示**：如果您尚未安装 JAX，建议先根据 [JAX 官方文档](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Finstallation.html) 完成 JAX 及其加速后端（CPU\u002FGPU\u002FTPU）的安装。\n\n## 安装步骤\n\n您可以选择从 PyPI 安装稳定版，或从 GitHub 安装最新开发版。\n\n### 方式一：安装稳定版（推荐）\n\n使用 pip 直接安装：\n\n```sh\npip install chex\n```\n\n**国内加速方案**：\n如果您在中国大陆地区，建议使用国内镜像源以加快下载速度：\n\n```sh\npip install chex -i https:\u002F\u002Fpypi.tuna.tsinghua.edu.cn\u002Fsimple\n```\n\n### 方式二：安装开发版\n\n如需体验最新功能或修复：\n\n```sh\npip install git+https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex.git\n```\n\n国内加速：\n\n```sh\npip install git+https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex.git -i https:\u002F\u002Fpypi.tuna.tsinghua.edu.cn\u002Fsimple\n```\n\n## 基本使用\n\nChex 的核心功能主要包括：**JAX 友好的数据类**、**运行时断言**以及**多变体测试**。以下是三个最常用的场景示例。\n\n### 1. 使用 JAX 兼容的 Dataclass\n\nChex 提供的 `@chex.dataclass` 装饰器可以将 Python 数据类注册为 JAX PyTree 节点，使其能够直接被 `jax.tree_util` 处理，同时也兼容 `dm-tree`。\n\n```python\nimport jax.numpy as jnp\nimport jax\nimport chex\n\n@chex.dataclass\nclass Parameters:\n  x: chex.ArrayDevice\n  y: chex.ArrayDevice\n\n# 初始化参数 (注意：必须使用关键字参数，不支持位置参数)\nparams = Parameters(\n    x=jnp.ones((2, 2)),\n    y=jnp.ones((1, 2)),\n)\n\n# 作为 JAX PyTree 进行操作\ndoubled_params = jax.tree_util.tree_map(lambda x: 2.0 * x, params)\n\n# 作为 Mapping (字典) 进行扁平化处理\nimport tree\nflat_list, treedef = tree.flatten(params)\n```\n\n### 2. 使用断言验证张量属性\n\n在调试 JAX 代码时，标准的类型注解无法检查 Shape、Rank 或 dtype。Chex 提供了一系列断言函数来在运行时验证这些属性。\n\n```python\nimport chex\nimport jax.numpy as jnp\n\nx = jnp.zeros((2, 3))\ny = jnp.zeros((2, 3))\nz = jnp.array([1.0, 2.0])\n\n# 验证形状是否一致\nchex.assert_equal_shape([x, y])\n\n# 验证特定形状\nchex.assert_shape(x, (2, 3))\n\n# 验证秩 (Rank)\nchex.assert_rank(z, 1) \n\n# 验证数值是否有限 (无 NaN 或 Inf)\nchex.assert_tree_all_finite({\"weights\": x, \"bias\": z})\n\n# 在 JIT 编译函数中使用断言 (需配合 chexify)\n@chex.chexify\n@jax.jit\ndef safe_compute(data):\n    chex.assert_tree_all_finite(data)\n    return jnp.sum(data)\n\nsafe_compute(x)\n```\n\n### 3. 使用 Variants 进行多模式测试\n\nJAX 代码在 `jit` 编译前后或在不同设备上的行为可能不同。Chex 的 `variants` 装饰器可以自动让同一个测试用例在多种配置下运行（例如：有\u002F无 JIT，有\u002F无设备内存）。\n\n```python\nimport chex\nimport jax.numpy as jnp\n\ndef my_function(x, y):\n    return x + y\n\nclass MyTest(chex.TestCase):\n\n  # 定义测试变体：同时测试 jit 编译版和非编译版\n  @chex.variants(with_jit=True, without_jit=True)\n  def test_addition(self):\n    # self.variant 会根据当前变体自动包装函数\n    tested_fn = self.variant(my_function)\n    \n    x = jnp.array([1.0, 2.0])\n    y = jnp.array([3.0, 4.0])\n    \n    result = tested_fn(x, y)\n    chex.assert_equal_shape([result, x])\n    chex.assert_trees_all_close(result, jnp.array([4.0, 6.0]))\n\n# 运行测试 (通常在命令行通过 pytest 或 python -m unittest 运行)\n# 此测试会自动运行两次：一次纯 Python，一次 JIT 编译\n```\n\n通过以上基础功能，您可以显著提升 JAX 项目的代码健壮性和测试覆盖率。更多高级用法（如 `assert_max_traces` 防止重复追踪）请参考官方 API 文档。","某深度学习团队正在基于 JAX 开发一套复杂的强化学习算法，需频繁处理多维状态张量与模型参数树。\n\n### 没有 chex 时\n- 调试形状错误极其耗时，开发者常因张量维度不匹配（如 `(32, 4)` 误传为 `(32,)`）导致程序在训练中途崩溃，却难以定位源头。\n- 缺乏对数据类型和秩（Rank）的自动校验，手动编写 `if x.shape != ...` 等断言代码冗长且容易遗漏，降低了代码可读性。\n- 自定义的参数结构无法直接兼容 JAX 的 `tree_map` 或 `dm-tree` 操作，必须反复进行字典转换，增加了样板代码和维护成本。\n- 难以统一测试代码在即时编译（jit）与非编译模式下的行为差异，导致某些 Bug 仅在部署时暴露。\n- 出现 NaN 或无穷大数值时，往往要跑完整个 epoch 才能发现，无法在计算发生的当下立即拦截。\n\n### 使用 chex 后\n- 利用 `assert_shape` 和 `assert_rank` 在函数入口处一键校验张量属性，形状不匹配时立即抛出清晰报错，将调试时间从小时级缩短至分钟级。\n- 通过 `assert_type` 和 `assert_equal_shape` 用声明式语法替代繁琐的手动检查，代码更简洁且逻辑意图一目了然。\n- 借助 `@chex.dataclass` 定义参数类，无需任何额外转换即可直接作为 JAX PyTree 进行映射操作，完美融合面向对象与函数式编程。\n- 使用变体测试工具轻松覆盖 jit 与非 jit 场景，确保逻辑在不同执行模式下表现一致，提前消除部署隐患。\n- 调用 `assert_tree_all_finite` 实时监控数值稳定性，一旦检测到非法数值立刻中断并定位，避免无效训练资源的浪费。\n\nchex 通过提供专为 JAX 设计的断言、数据结构与测试工具，将原本脆弱的数值计算流程转变为健壮、可信赖的工程系统。","https:\u002F\u002Foss.gittoolsai.com\u002Fimages\u002Fgoogle-deepmind_chex_70d56cf8.png","google-deepmind","Google DeepMind","https:\u002F\u002Foss.gittoolsai.com\u002Favatars\u002Fgoogle-deepmind_06b1dd17.png","","https:\u002F\u002Fwww.deepmind.com\u002F","https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind",[80,84],{"name":81,"color":82,"percentage":83},"Python","#3572A5",98.9,{"name":85,"color":86,"percentage":87},"Shell","#89e051",1.1,936,68,"2026-04-18T15:12:12","Apache-2.0",1,"未说明","非必需。库提供断言函数检查 GPU\u002FTPU 可用性，但运行库本身不强制要求特定显卡型号、显存大小或 CUDA 版本。",{"notes":96,"python":97,"dependencies":98},"Chex 是用于编写可靠 JAX 代码的实用工具库，而非独立的深度学习模型。它主要用于代码检测、调试（如在上下文管理器中将 pmap 转换为 vmap）以及测试不同变体（如 JIT 编译与非编译版本）。安装可通过 PyPI (pip install chex) 或 GitHub 源码进行。使用时需注意：Chex 的数据类不支持位置参数构造；在 JIT 编译函数中使用值断言需配合 @chex.chexify 装饰器；测试变体功能需继承 chex.TestCase。","3.7+",[99,100,101],"jax","dm-tree","absl-py",[14],"2026-03-27T02:49:30.150509","2026-04-20T10:24:22.153565",[106,111,116,121,126,131],{"id":107,"question_zh":108,"answer_zh":109,"source_url":110},44560,"Conda 安装的 Chex 包依赖项中的 JAX 版本要求有误怎么办？","这是 conda-forge 渠道上的一个已知问题，特定版本（如 0.1.7）的依赖项中 JAX 版本要求写错了（例如写成了 >=0.1.55，实际应为 >=0.4.6）。维护者通常会修复此问题并推送新版本。如果遇到此问题，请等待几分钟让 conda-forge 更新，或者手动指定正确的 JAX 版本进行安装。","https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fissues\u002F270",{"id":112,"question_zh":113,"answer_zh":114,"source_url":115},44561,"导入 Chex 时出现 'ModuleNotFoundError: No module named jax.numpy' 错误如何解决？","该错误通常发生在 JAX 版本过低（低于 0.4.1）或环境安装不完整时。`import jax.numpy` 语法在旧版 JAX 中可能不被支持。解决方案是升级 JAX 到最新版本（建议 >=0.4.1），或者检查当前环境是否损坏并重新安装。确认 Chex 在 jax==0.3.14 及以上版本中通常能正常工作，但推荐使用更新的稳定版。","https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fissues\u002F226",{"id":117,"question_zh":118,"answer_zh":119,"source_url":120},44562,"遇到 'AttributeError: module jax has no attribute _src' 错误是什么原因？","这通常是由于 JAX、Chex 和 Optax 之间的版本不兼容导致的。当 JAX 版本更新（如 >0.3.17）而 Optax 或 Chex 版本过旧时，可能会尝试访问已被移除或移动的私有属性（如 `_src`）。解决方法是确保同时升级 optax、chex 和 jax 到相互兼容的最新版本。","https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fissues\u002F197",{"id":122,"question_zh":123,"answer_zh":124,"source_url":125},44563,"升级 JAX 后出现 'CpuDevice no longer in jax' 错误如何处理？","这是因为新版 JAX（如 0.3.7+）移除了 `CpuDevice` 类等内部实现细节。如果使用的是旧版 Chex（如 0.1.3 之前），会因引用这些被移除的类而报错。解决方案是将 Chex 升级到 v0.1.3 或更高版本，该版本已包含对此变化的修复。","https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fissues\u002F156",{"id":127,"question_zh":128,"answer_zh":129,"source_url":130},44564,"为什么元组（tuple）不被识别为有效的 chex.ArrayTree 类型？","`chex.ArrayTree` 类型定义明确排除了 Python 原生标量类型（如 float, int），仅允许 `numpy.number` 或 `numpy.bool_` 类型的标量。此外，使用的类型检查器（如 Pytype）对异构列表（list）的检查较宽松，而对元组（tuple）的检查更严格，导致包含原生浮点数的元组无法通过类型检查。若需更通用的类型定义，目前 Chex 尚未提供 `chex.Pytree`，建议尽量使用 numpy 标量或列表结构替代。","https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fissues\u002F323",{"id":132,"question_zh":133,"answer_zh":134,"source_url":135},44565,"运行测试套件时出现 'AssertsChexifyTest.test_uninspected_checks' 失败怎么办？","这通常是测试环境配置问题或依赖版本不匹配导致的特定测试用例失败。虽然日志显示了大量跳过（s）和通过（.）的测试，但该特定失败可能与内部断言逻辑有关。建议首先确保所有依赖（jax, jaxlib, pytest 等）均为最新兼容版本。如果问题依旧，可能是特定平台或 Python 版本的偶发问题，可尝试忽略该特定测试或查看 CI 流水线上的通过情况作为参考。","https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fissues\u002F204",[137,142,147,152,157,162,167,172,177,182,187,192,197,202,207,212,217,222,227,232],{"id":138,"version":139,"summary_zh":140,"released_at":141},351978,"v0.1.91","## 变更内容\n* 将对 `jax.lib.xla_extension` 的引用迁移到 `jax.extend`。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F403 中完成。\n* 添加 `assert_tree_shape`，它类似于 `assert_shape`，但适用于树结构中的所有叶节点。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F404 中完成。\n* 在构造 `assert_tree_shape` 错误信息时使用 `format_shape_matcher`。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F405 中完成。\n* 迁移到 `pyproject.toml`。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F406 中完成。\n* 发布 v0.1.91 版本。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F407 中完成。\n\n\n**完整变更日志**: https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fcompare\u002Fv0.1.90...v0.1.91","2025-09-01T21:48:43",{"id":143,"version":144,"summary_zh":145,"released_at":146},351979,"v0.1.90","## 变更内容\n* 修复源代码链接，使其指向 master 分支。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F384 中完成。\n* chex：避免使用已弃用的 `jax.util.unzip2`。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F391 中完成。\n* 更新 JAX 夜间构建索引的使用方式。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F392 中完成。\n* [chex] 移除过时的 try\u002Fexcept 导入语句。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F396 中完成。\n* 在 https:\u002F\u002Fgithub.com\u002Fjax-ml\u002Fjax\u002Fcommit\u002F06448864abd6e8187e5b4d9b1ff08ab14fe3b8e0 之后更新 restrict_backends。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F397 中完成。\n* 如果可用，则使用 backend_compile_and_load；否则使用 backend_compile。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F398 中完成。\n* 在 assert_trees_all_close_ulp 中支持 BF16 数据类型。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F399 中完成。\n* 修复拼写错误。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F400 中完成。\n* 发布 v0.1.90 版本。由 @copybara-service[bot] 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F402 中完成。\n\n\n**完整变更日志**：https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fcompare\u002Fv0.1.89...v0.1.90","2025-07-23T19:50:19",{"id":148,"version":149,"summary_zh":150,"released_at":151},351980,"v0.1.89","## 变更内容\n* 切换到新的 `unittest` 测试加载器 API，由 @daskol 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F372 中完成\n* 添加针对 JAX 夜间版的测试，由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F376 中完成\n* CI：测试 Python 3.12，由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F377 中完成\n* Chex：添加 Python 3.13 的 CI 测试，由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F378 中完成\n* 修复拼写错误：“assetion” -> “assertion”，由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F379 中完成\n* 维护与 Python 3.13 的兼容性，由 @daskol 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F375 中完成\n\n## 新贡献者\n* @daskol 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F372 中完成了首次贡献\n\n**完整变更日志**：https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fcompare\u002Fv0.1.88...v0.1.89","2025-02-25T15:42:58",{"id":153,"version":154,"summary_zh":155,"released_at":156},351981,"v0.1.88","## 变更内容\n* 由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F366 中修复了 `test_assert_trees_all_equal_prng_keys`\n* 避免依赖 jax.lib.xla_bridge._backends，因为该模块为私有且即将被删除，由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F367 中完成\n* 发布版本 0.1.88，由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F369 中完成\n\n\n**完整变更日志**: https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fcompare\u002Fv0.1.87...v0.1.88","2024-12-11T10:33:44",{"id":158,"version":159,"summary_zh":160,"released_at":161},351996,"v0.1.1","## What's Changed\r\n* Move dataclass registration to __init__ so that it's invoked after deserialization. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F111\r\n* Add pytype for jax array's dtypes. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F112\r\n* Fix dataclass registration on deserialization. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F114\r\n* Fix restrict_backends after jax.xla.backend_compile was moved by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F116\r\n* Refactor asserts.py and warn users not to rely on asserts_internal's functionality. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F117\r\n* Set up ReadTheDoc pages and add a few examples. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F118\r\n* Include Sphinx builds into CI tests. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F119\r\n* Adds internal functionality by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F122\r\n* Update Chex citation. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F125\r\n* Refactor assertions in preparation for including them into the RTD docs. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F126\r\n* Add asserts, variants, and pytypes modules to the RTD docs. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F127\r\n* Fix references to collections.abc.Mappable -> collections.abc.Mapping in docs and comments. collections.abc.Mappable does not exist. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F129\r\n* Document the rational behing the mappability of chex.dataclasses. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F130\r\n* Add 3 new tree assertions: by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F131\r\n* Add `assert_tree_is_sharded` for asserting that a tree is sharded across the specified devices. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F132\r\n* Add PyTreeDef to pytypes.  by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F134\r\n* Disallow `ShardedDeviceArray`s in `assert_tree_is_on_host` and `assert_tree_is_on_device`. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F133\r\n* Bump ipython from 7.16.1 to 7.16.3 in \u002Frequirements by @dependabot in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F135\r\n* Remove the old venv directory before testing the package. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F138\r\n* Refactor asserts.py in preparation for experimental device assertions. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F137\r\n* Fix minor typo in docs. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F139\r\n* Improve exception message for `assert_tree_shape_prefix`. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F143\r\n* Release v0.1.1 by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F146\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fcompare\u002Fv0.1.0...v0.1.1","2022-02-25T12:54:43",{"id":163,"version":164,"summary_zh":165,"released_at":166},351997,"v0.1.0","Support for Python 3.6 has been dropped as per [JAX deprecation policy](https:\u002F\u002Fjax.readthedocs.io\u002Fen\u002Flatest\u002Fdeprecation.html). Please upgrade to a supported Python version.","2021-11-18T10:06:13",{"id":168,"version":169,"summary_zh":170,"released_at":171},351992,"v0.1.5","## What's Changed\r\n* Add support for value assertions in jitted functions. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F178\r\n* [JAX] Avoid private implementation detail _ScalarMeta. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F180\r\n* [JAX] Avoid implicit references to jax._src. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F181\r\n* Release v0.1.15 by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F184\r\n\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fcompare\u002Fv0.1.4...v0.1.5","2022-09-13T09:55:28",{"id":173,"version":174,"summary_zh":175,"released_at":176},351993,"v0.1.4","## What's Changed\r\n* Add an InitVar field in the dataclass tests. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F161\r\n* Download latest .pylintrc version in tests. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F167\r\n* Fix `assert_axis_dimension_comparator` usages. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F168\r\n* Update \"jax.tree_util\" functions by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F171\r\n* Use jax.tree_util.tree_map in place of deprecated tree_multimap. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F175\r\n* Silence some pytype errors. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F174\r\n* Add chex.Dimensions utility for readable shape asserts. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F169\r\n\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fcompare\u002Fv0.1.3...v0.1.4","2022-08-04T06:58:28",{"id":178,"version":179,"summary_zh":180,"released_at":181},351994,"v0.1.3","## What's Changed\r\n* Slight helping clarification to `clear_trace_counter`. by @lucasb-eyer in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F148\r\n* Add new JAX-specific pytypes to chex pytypes. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F153\r\n* Remove `chex.{C,G,T}puDevice` in favour of `chex.Device`. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F154\r\n\r\n## New Contributors\r\n* @lucasb-eyer made their first contribution in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F148\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fcompare\u002Fv0.1.2...v0.1.3","2022-04-19T09:42:46",{"id":183,"version":184,"summary_zh":185,"released_at":186},351995,"v0.1.2","## What's Changed\r\n* Support JAX parallel operations in chex.fake_pmap contexts by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F142\r\n* Remove references to jax.numpy.lax_numpy. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F150\r\n\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fcompare\u002Fv0.1.1...v0.1.2","2022-03-31T17:04:45",{"id":188,"version":189,"summary_zh":190,"released_at":191},351982,"v0.1.87","## 变更内容\n* 在 chex.Dimensions 中添加展平功能。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F341 中完成。\n* 修复 warn_deprecated_function 中的文档字符串。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F343 中完成。\n* 添加用于创建已弃用别名的实用工具。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F345 中完成。\n* 修复 linter 的警告。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F349 中完成。\n* 测试 RTD 网页钩子。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F350 中完成。\n* 移除一个有问题的异步 chexify 测试，因为以下情况无法保证：由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F353 中完成。\n* 将已弃用的 `jax.tree_*` 函数替换为 `jax.tree.*`。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F354 中完成。\n* 禁用在 __main__ 中定义的数据类在 JAX 树工具中的注册。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F356 中完成。\n* Chex：为支持的最旧 JAX 版本添加 CI 作业。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F358 中完成。\n* 将 warn_deprecated_function 中的 `replacement` 关键字参数设为可选。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F362 中完成。\n* 为 PyPI 发布提高版本号。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F363 中完成。\n\n\n**完整变更日志**：https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fcompare\u002Fv0.1.86...v0.1.87","2024-09-30T14:40:29",{"id":193,"version":194,"summary_zh":195,"released_at":196},351983,"v0.1.86","## 变更内容\n* 避免使用已弃用的 `jax.Array` 的 `device_buffer` 属性，由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F325 中完成  \n* 更新 `test.sh` 脚本，由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F331 中完成  \n* 更新 `chex.assert_type`，使其检查具体类型，而不仅仅是断言类型为浮点或整数子类型。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F333 中完成  \n* 调整 `test_assert_trees_all_close` 的测试容差，由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F334 中完成  \n* 在定义时注册数据类，以便在大多数情况下（除一种例外）都能成功反序列化数据类对象、类型和树形结构定义。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F336 中完成  \n* 在 `chex` 中添加警告信息。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F337 中完成  \n* 添加 `Dimensions.size` 方法，用于从命名形状获取扁平化的大小。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F339 中完成  \n* 提升版本号以发布新版本。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F342 中完成  \n\n\n**完整变更日志**: https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fcompare\u002Fv0.1.85...v0.1.86","2024-03-19T10:05:05",{"id":198,"version":199,"summary_zh":200,"released_at":201},351984,"v0.1.85","## 变更内容\n* [JAX] 将 `jax.devices(\"cpu\")` 的用法替换为 `jax.local_devices(backend=\"cpu\")`。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fpull\u002F317 中完成。\n\n\n**完整变更日志**: https:\u002F\u002Fgithub.com\u002Fgoogle-deepmind\u002Fchex\u002Fcompare\u002Fv0.1.84...v0.1.85","2023-11-22T18:26:22",{"id":203,"version":204,"summary_zh":205,"released_at":206},351985,"v0.1.84","更新 `pytype.py` 中的类型定义，并修复各类警告。","2023-10-20T09:58:06",{"id":208,"version":209,"summary_zh":210,"released_at":211},351986,"v0.1.83","* 移除已弃用的 `assert_tree_*` 方法；\r\n* 提升 dataclass 的易用性；\r\n* 根据 Jax 的最新变更，更新依赖项和 pytypes。","2023-09-20T15:35:03",{"id":213,"version":214,"summary_zh":215,"released_at":216},351987,"v0.1.82","## 变更内容\nChex 已从外部的 `dm-tree` 迁移到 JAX 原生的 `jax.tree_util` 库来操作 `PyTree`。这意味着，**`None` 将不再被 Chex 的树断言视为独立的叶节点**。\n需要注意的是，您仍然可以使用 `chex.assert_tree_no_nones()` 来检查是否存在 `None`。\n\n* 移除断言中对 dm-tree 的使用。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F288 中完成。\n* 为 chex.dataclass 添加对 `kw_only` 的支持。由 @jblespiau 在 https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F290 中完成。\n* 弃用 `ignore_nones`。由 @copybara-service 在 https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F292 中完成。\n\n\n**完整变更日志**: https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fcompare\u002Fv0.1.81...v0.1.82","2023-07-20T10:23:48",{"id":218,"version":219,"summary_zh":220,"released_at":221},351988,"v0.1.81","Update requirements.txt to use `numpy>=1.25.0`.","2023-06-29T10:59:08",{"id":223,"version":224,"summary_zh":225,"released_at":226},351989,"v0.1.8","## What's Changed\r\n* Drop support for python 3.8, as per JAX deprecation policy.\r\n* Use new `jax.tree_util.register_pytree_with_keys` to register `chex.Dataclass`. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F266\r\n* Add size based assertions by @tttc3 in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F265\r\n* Integer sequence support for `chex.assert_tree_shape{_prefix, _suffix}` by @hylkedonker in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F275\r\n* Resolve unsoundness caught by pytype --strict-none-binding. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F276\r\n* Release v0.1.8. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F281\r\n\r\n## New Contributors\r\n* @tttc3 made their first contribution in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F265\r\n* @hylkedonker made their first contribution in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F275\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fcompare\u002Fv0.1.7...v0.1.8","2023-06-27T13:14:51",{"id":228,"version":229,"summary_zh":230,"released_at":231},351990,"v0.1.7","## What's Changed\r\n* Bump the minimum Jax version to 0.4.6 \r\n* feat(ci): bump `setup-python` version and enable cache by @SauravMaheshkar in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F239\r\n* [README] fix link to `@mappable_dataclass` by @pfackeldey in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F241\r\n* Add ArrayDeviceTree and ArrayNumpyTree pytypes + improve ArrayTree. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F253\r\n* Release v0.1.7 by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F256\r\n\r\n## New Contributors\r\n* @SauravMaheshkar made their first contribution in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F239\r\n* @pfackeldey made their first contribution in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F241\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fcompare\u002Fv0.1.6...v0.1.7","2023-03-27T10:43:40",{"id":233,"version":234,"summary_zh":235,"released_at":236},351991,"v0.1.6","## What's Changed\r\n* Replace `jnp.DeviceArray` with the new public type `jax.Array`. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F187\r\n* Unpack assertions' kwargs at a higher level for improved readability. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F209\r\n* Allow calling `@chexify`-ed functions from non-main threads. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F214\r\n* Add `assert_trees_all_equal_shapes_and_dtypes` by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F219\r\n* Implement a new unflatten for chex.dataclass that avoids __init__ while keeping the (un)flattened order unchanged. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F222\r\n* Fix examples in the docs & linter warnings. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F235\r\n* Replace uses of jax.lib.xla_extension.Device with jax.Device. by @copybara-service in https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fpull\u002F237\r\n\r\n\r\n**Full Changelog**: https:\u002F\u002Fgithub.com\u002Fdeepmind\u002Fchex\u002Fcompare\u002Fv0.1.5...v0.1.6","2023-02-06T12:09:21"]