Flax traverse_util 全面指南嵌套字典的扁平化、路径感知映射与不可变数据遍历【免费下载链接】flaxFlax is a neural network library for JAX that is designed for flexibility.项目地址: https://gitcode.com/GitHub_Trending/fl/flaxFlax 的flax.traverse_util模块为 JAX 生态中的嵌套数据结构尤其是模型参数、梯度、状态这类嵌套字典提供了一套小巧而强大的工具既能将嵌套字典扁平化为元组键字典flatten_dict/unflatten_dict也能在保留路径信息的前提下逐叶子映射path_aware_map还提供了一套可组合的不可变 Traversal 遍历 API。阅读本文后你将掌握参数树的扁平化/还原、按路径条件化处理参数如搭配 optax 的multi_transform实现分层学习率、以及用 Traversal 在不修改原数据的前提下精准选中并更新任意子结构的能力这些技巧在模型手术、checkpoint 处理与优化器配置中会反复用到。本文以 docs_nnx/api_reference/flax.traverse_util.rst 的 API 文档为骨架结合 flax/traverse_util.py 的完整源码实现与 tests/traverse_util_test.py 的测试用例进行纵深讲解。模块定位为什么需要遍历与扁平化工具JAX 的核心抽象是 pytree任意嵌套的 dict、list、tuple 等容器都可以被递归遍历叶子通常是jax.Array或标量。模型参数、优化器状态、梯度在 Flax 中几乎都是嵌套字典VariableDict / FrozenDict形态。然而实际工程中经常遇到两类需求整体拍平把任意深度的嵌套字典变成扁平字典键为路径元组方便按路径精确检索、排序、序列化或批量处理定向更新只对树中符合某种路径条件的子结构做变换且绝不原地修改——这是 JAX 函数式风格的核心约束。flax.traverse_util正是为解决这两类问题而存在。模块 docstring 开宗明义A utility for traversing immutable datastructures遍历不可变数据结构的工具并强调Traversals never mutate the original dataTraversal 永不修改原始数据一次update本质上是返回一份包含指定更新的数据副本。该模块位于 flax/traverse_util.py被 flax/training/checkpoints.py、flax/linen/module.py 等核心文件直接依赖是整个 Flax 基础设施的底层件之一。一、flatten_dict把嵌套字典拍平为路径键字典flatten_dict是模块中最常用的函数负责把嵌套字典展平。其 API 文档见 docs_nnx/api_reference/flax.traverse_util.rst 的Dict utils一节与 docstring 给出的核心示例如下from flax.traverse_util import flatten_dict xs {foo: 1, bar: {a: 2, b: {}}} flat_xs flatten_dict(xs) # flat_xs # {(foo,): 1, (bar, a): 2}注意空字典默认被忽略{bar: {b: {}}}中的空字典b不会出现在结果里unflatten_dict也无法还原它。参数详解源码签名flax/traverse_util.pydef flatten_dict(xs, keep_empty_nodesFalse, is_leafNone, sepNone):参数默认值作用xs必填输入的嵌套字典必须是dict或flax.core.FrozenDict否则触发assert断言失败源码第 137-139 行keep_empty_nodesFalse为True时空字典不再被丢弃而是以模块级哨兵对象traverse_util.empty_node作为值保留便于无损往返is_leafNone可选函数接收(prefix, xs)两个参数返回True表示当前嵌套字典应被视为叶子、不再继续展开sepNone若指定如/返回字典的键由路径元组改为sep连接而成的字符串None时键为元组keep_empty_nodes与empty_node哨兵源码中empty_node被定义为struct.dataclass class _EmptyNode: pass empty_node _EmptyNode()它特意用flax.struct.dataclass装饰注释明确说明是为了 be compatible with JAX与 JAX 兼容这样哨兵可以作为 pytree 叶子参与jax.jit等变换。开启keep_empty_nodesTrue后flat_xs flatten_dict(xs, keep_empty_nodesTrue) # {(foo,): 1, (bar, a): 2, (bar, b): traverse_util.empty_node} xs_restore unflatten_dict(flat_xs) # {foo: 1, bar: {a: 2, b: {}}} —— 与原始输入完全一致测试 tests/traverse_util_test.py 的test_flatten_dict_keep_empty验证了这一往返一致性。is_leaf自定义叶子判定当你想把某个深度的子字典整体当作叶子值保留时使用is_leaf。测试test_flatten_dict_is_leaftests/traverse_util_test.py展示了典型用法xs {foo: {c: 4}, bar: {a: 2, b: {}}} flat_xs flatten_dict( xs, is_leaflambda k, x: len(k) 1 and len(x) 2 ) # {(foo, c): 4, (bar,): {a: 2, b: {}}} # —— bar 因满足 len(路径)1 且 len(字典)2 而被视为叶子整体保留注意is_leaf的判定发生在递归展开之前当某节点既满足叶子条件又仍是 dict 时它整体作为一个值被保留其内部结构不再展开。sep路径字符串化测试test_flatten_dicttests/traverse_util_test.py验证了sep/的用法flat_xs flatten_dict(xs, sep/) # {foo: 1, bar/a: 2}这在需要把路径直接拼进文件名、日志标签或 GCS/tensorstore 子路径的场景中非常实用下文 checkpoint 部分会看到真实案例。二、unflatten_dict还原嵌套结构unflatten_dict是flatten_dict的逆操作flax/traverse_util.pyfrom flax.traverse_util import unflatten_dict flat_xs { (foo,): 1, (bar, a): 2, } xs unflatten_dict(flat_xs) # {foo: 1, bar: {a: 2}}要点输入必须是普通dict键为路径元组或字符串配合sep使用若值与empty_node相等还原时自动替换为{}对应keep_empty_nodesTrue的往返中途路径缺失时按需自动创建中间层字典与flatten_dict使用同一个sep参数路径字符串会按sep切分回元组。源码第 165-178 行展示了实现遍历每个(path, value)沿路径逐层下沉创建cursor字典最后把值挂到叶子键上。三、path_aware_map带路径信息的叶子级映射path_aware_map是路径感知的 mapflax/traverse_util.py它对嵌套字典的每个叶子调用f(path, value)其中path是该叶子从根到自身的路径元组。docstring 示例import jax.numpy as jnp from flax import traverse_util params {a: {x: 10, y: 3}, b: {x: 20}} f lambda path, x: x 5 if x in path else -x traverse_util.path_aware_map(f, params) # {a: {x: 15, y: -3}, b: {x: 25}}实现原理源码实现只有两行核心逻辑flat flatten_dict(nested_dict, keep_empty_nodesTrue) return unflatten_dict( {k: f(k, v) if v is not empty_node else v for k, v in flat.items()} )即先无损拍平keep_empty_nodesTrue保住空节点对每个叶子应用f再把结果还原为嵌套结构。因为flatten_dict/unflatten_dict都是纯函数式操作path_aware_map天然不修改输入并能在结果中保留空字典测试test_path_aware_map_with_empty_nodestests/traverse_util_test.py。实战与 optaxmulti_transform配合实现分层优化这是path_aware_map最典型的落地场景——为不同路径的参数贴上不同标签交给 optax 做多组优化器调度。测试test_path_aware_map_with_multi_transformtests/traverse_util_test.py给出了完整可运行示例params { linear_1: {w: jnp.zeros((5, 6)), b: jnp.zeros(5)}, linear_2: {w: jnp.zeros((6, 1)), b: jnp.zeros(1)}, } gradients jax.tree_util.tree_map(jnp.ones_like, params) # 占位梯度 # 按路径是否为 w 打标签kernel / bias param_labels traverse_util.path_aware_map( lambda path, x: kernel if w in path else bias, params ) tx optax.multi_transform( {kernel: optax.sgd(1.0), bias: optax.set_to_zero()}, param_labels ) state tx.init(params) updates, new_state tx.update(gradients, state, params) new_params optax.apply_updates(params, updates)效果w权重使用 SGD 更新b偏置完全冻结——最终b与原始值完全一致、w被更新。同样的模式也适用于optax.masked测试test_path_aware_map_with_maskedtests/traverse_util_test.py可见该函数是按参数名/路径做选择性优化的通用前置工具。四、Traversal 遍历 API可组合的不可变数据结构访问器模块 docstring 与源码共同定义了一套面向对象风格的 Traversal 体系。它的设计哲学是Traversal 是一个透镜lens——它选中数据结构中的一个子集支持iterate读取与update写入副本并可通过组合构建出任意复杂的选区。4.1 基础用法选中与更新从 identity 遍历t_identity出发通过方法链和属性/下标访问构造目标选区from flax import traverse_util import dataclasses dataclasses.dataclass class Foo: foo: int 0 bar: int 0 # 属性遍历选中 Foo 的 foo 属性 x Foo(foo1) iterator traverse_util.TraverseAttr(foo).iterate(x) list(iterator) # [1] # 组合遍历每个元素取其 foo 键 data [{foo: 1, bar: 2}, {foo: 3, bar: 4}] traversal traverse_util.t_identity.each()[foo] list(traversal.iterate(data)) # [1, 3] # update不修改原对象返回更新后的副本 data {foo: Foo(bar2)} traversal traverse_util.t_identity[foo].bar data traversal.update(lambda x: x x, data) # {foo: Foo(foo0, bar4)}4.2 Traversal 类族与组合原语核心抽象类Traversalflax/traverse_util.py定义了两个抽象方法并提供了组合方法方法作用update(fn, inputs)对选中的每个元素应用fn返回新对象不可变风格iterate(inputs)返回一个迭代器产出选中的元素set(values, inputs)用一组新值覆盖选区值数量不匹配时抛ValueErrorcompose(other)组合两个 Traversal先走外层、再走内层merge(*traversals)合并多个 Traversal 的选区用于同时选中多个分支each()遍历容器中的每个元素dict/list/tupletree()遍历 pytree 的每个叶子filter(fn)按谓词过滤选中的值__getattr__/__getitem__语法糖.attr等价于compose(TraverseAttr(attr))[key]等价于compose(TraverseItem(key))具体实现类包括TraverseId恒等遍历iterate产出自身update直接应用fn全局单例t_identity是它源码第 296-308 行TraverseAttr遍历对象的属性支持 namedtuple_replace、dataclassdataclasses.replace及普通对象copy.copy后setattrTraverseItem遍历容器下标或键支持元组、namedtuple、list、dict 以及slice 切片t_identity[1:3]可一次选中多个元组元素TraverseEach遍历 list/tuple/dict 中每个条目对非这三种类型抛ValueErrorTraverseTree基于jax.tree_util遍历 pytree 所有叶子TraverseFilter按谓词过滤TraverseMerge将多个 Traversal 的选区合并成一个选区TraverseCompose串联两个 Traversal是上述所有组合方法的底层引擎。4.3 测试覆盖的典型组合tests/traverse_util_test.py 的TraversalTest类逐条验证了这些行为# 元组下标与切片 x (1, 2, 3, 4) list(traverse_util.t_identity[1:3].iterate(x)) # [2, 3] traverse_util.t_identity[1:3].update(lambda x: x x, x) # (1, 4, 6, 4) # namedtuple 属性 Point collections.namedtuple(Point, [x, y]) x Point(x1, y2) traverse_util.t_identity.y.update(lambda x: x x, x) # Point(x1, y4) # each merge同时处理多个字段 x [{foo: 1, bar: 2}, {foo: 3, bar: 4}] t traverse_util.t_identity.each().merge( traverse_util.TraverseItem(foo), traverse_util.TraverseItem(bar) ) list(t.iterate(x)) # [1, 2, 3, 4] # filter按条件过滤后更新 x [1, -2, 3, -4] t traverse_util.t_identity.each().filter(lambda x: x 0) t.update(lambda x: -x, x) # [1, 2, 3, 4] # set校验值数量 t traverse_util.t_identity[foo].each() t.set([3, 4], {foo: [1, 2]}) # {foo: [3, 4]} # 值太少 / 太多都会抛 ValueError4.4 弃用提示Traversal 与flax.optim源码第 214-223 行在Traversal.__new__中发出DeprecationWarning提示flax.traverse_util.Traversal将被弃用。如果你是为了flax.optim使用它请改用optax详细迁移说明见官方 optax 更新指南docs/guides/converting_and_upgrading/optax_update_guide.rst 可找到对应的迁移思路。也就是说新代码中处理优化器相关需求请优先使用 optax 及上文介绍的path_aware_map方案Traversal 类目前仍保留t_identity定义时用warnings.catch_warnings()抑制了弃用告警保证其可用主要服务于旧的flax.optim兼容路径与更广义的不可变结构遍历需求。4.5ModelParamTraversal按参数全名筛选ModelParamTraversalflax/traverse_util.py专为模型参数设计构造时传入filter_fn它会收到形如/module/sub_module/parameter_name的完整路径字符串与参数值返回该参数是否被选中。测试test_param_selectiontests/traverse_util_test.py展示了按名称包含kernel筛选并翻倍更新的效果params {x: {kernel: 1, bias: 2, y: {kernel: 3, bias: 4}, z: {}}} traversal traverse_util.ModelParamTraversal( lambda name, _: kernel in name ) list(traversal.iterate(params)) # [1, 3]只选中两个 kernel traversal.update(lambda x: x x, params) # kernel 变为 2、6bias 不变它内部通过flatten_dict含keep_empty_nodesTrue获取扁平参数表按字典序排序后处理对flax.core.FrozenDict输入会返回FrozenDict结果保持不可变类型不变。同时它只接受嵌套 dict 或FrozenDict对其他类型抛ValueError测试test_only_works_on_model_params验证了这一点。五、仓库内的真实应用checkpoint 与 Linen 内部的依赖traverse_util并非孤立工具它在 Flax 基础设施中承担关键角色以下用法均有源码可查5.1 checkpoint 中的多进程数组MPA分片处理flax/training/checkpoints.py 的_split_mp_arrays在保存 checkpoint 时用flatten_dict(target, keep_empty_nodesTrue)拍平整个目标树挑出所有多进程分布式数组把路径/.join(key)拼成子路径并替换为占位符最后用unflatten_dict还原flattened traverse_util.flatten_dict(target, keep_empty_nodesTrue) mpa_targets [] for key, value in flattened.items(): if _is_multiprocess_array(value): subpath /.join(key) mpa_targets.append((value, subpath)) flattened[key] MP_ARRAY_PH subpath target traverse_util.unflatten_dict(flattened)恢复路径_restore_mpasflax/training/checkpoints.py同样依赖flatten_dict/unflatten_dict完成占位符替换与还原。这正是sep参数与keep_empty_nodes在真实工程中的价值体现。5.2 Linen 模块中对叶子或整个 dict 统一映射flax/linen/module.py 在map_相关逻辑中用flatten_dict普通版与keep_empty_nodesTrue版处理叶子值与嵌套字典混合的输入再以unflatten_dict恢复结构。flax.core.lift的 remat/scan 内部flax/core/lift.py也用它做中间表示的拍平与还原。这些使用点证明了该模块是 Flax 参数树序列化与变换管线的标准工具理解它能帮助你读懂 checkpoint 源码乃至自定义序列化逻辑。六、快速参考API 一览与选择建议API一句话定位典型场景flatten_dict(xs, keep_empty_nodes, is_leaf, sep)嵌套字典 → 路径键元组或字符串字典参数树拍平、按路径排序/检索、路径拼接unflatten_dict(xs, sepNone)路径键字典 → 嵌套字典还原拍平结果、checkpoint 占位符回填path_aware_map(f, nested_dict)对每个叶子调用f(path, value)并还原参数打标签喂给optax.multi_transform/optax.masked、按路径做条件变换t_identity/ Traversal 族可组合的不可变选区透镜选中并更新嵌套结构中的子集注意弃用提示ModelParamTraversal(filter_fn)按/module/param全名筛选参数传统flax.optim多优化器场景的参数分组选型建议新项目处理按路径映射优先用path_aware_map纯函数、无弃用风险、与 optax 天然契合需要整体扁平化做序列化/哈希/排序时用flatten_dictunflatten_dict注意空字典与keep_empty_nodes的取舍Traversal 类族适用于需要精确读写对象属性/切片/多分支合并的不可变结构操作但请知晓其DeprecationWarning背景。七、易错点与工程注意事项flatten_dict只接受 dict / FrozenDict传入 list、tuple 或其他容器会直接触发assert失败源码 137-139 行这类结构请先用jax.tree_util或在外面包一层 dict。空字典默认丢失若你的结构中有空 dict 且需要无损往返务必设置keep_empty_nodesTrue否则unflatten_dict无法恢复原状。unflatten_dict输入必须是普通 dict传入 FrozenDict 会触发断言如需保持 FrozenDict 类型可参考ModelParamTraversal的处理方式返回时手动包回FrozenDict。path_aware_map中f的签名是(path, value)与jax.tree_util.tree_map的单参数f(value)不同path是字符串元组可用于判断路径是否包含某个字段名。Traversal 的不可变性update返回新对象而不改原对象但需要注意TraverseItem对 list/dict 使用copy.copy的浅拷贝语义——叶子元素被替换但未被选中的嵌套子结构仍与原对象共享引用。版本背景本文描述的行为以当前仓库 flax/traverse_util.pyCopyright 2024为准Traversal 相关 API 处于弃用过渡期新代码应避开依赖弃用告警的调用方式。通过上述六个章节你已掌握flax.traverse_util的全部核心 API、底层实现机理、测试验证路径以及在 checkpoint 与 Linen 内部的真实工程用法——这套工具虽然轻量却是理解 Flax 参数树操作流水线的重要基石。【免费下载链接】flaxFlax is a neural network library for JAX that is designed for flexibility.项目地址: https://gitcode.com/GitHub_Trending/fl/flax创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Flax traverse_util 全面指南：嵌套字典的扁平化、路径感知映射与不可变数据遍历
Flax traverse_util 全面指南嵌套字典的扁平化、路径感知映射与不可变数据遍历【免费下载链接】flaxFlax is a neural network library for JAX that is designed for flexibility.项目地址: https://gitcode.com/GitHub_Trending/fl/flaxFlax 的flax.traverse_util模块为 JAX 生态中的嵌套数据结构尤其是模型参数、梯度、状态这类嵌套字典提供了一套小巧而强大的工具既能将嵌套字典扁平化为元组键字典flatten_dict/unflatten_dict也能在保留路径信息的前提下逐叶子映射path_aware_map还提供了一套可组合的不可变 Traversal 遍历 API。阅读本文后你将掌握参数树的扁平化/还原、按路径条件化处理参数如搭配 optax 的multi_transform实现分层学习率、以及用 Traversal 在不修改原数据的前提下精准选中并更新任意子结构的能力这些技巧在模型手术、checkpoint 处理与优化器配置中会反复用到。本文以 docs_nnx/api_reference/flax.traverse_util.rst 的 API 文档为骨架结合 flax/traverse_util.py 的完整源码实现与 tests/traverse_util_test.py 的测试用例进行纵深讲解。模块定位为什么需要遍历与扁平化工具JAX 的核心抽象是 pytree任意嵌套的 dict、list、tuple 等容器都可以被递归遍历叶子通常是jax.Array或标量。模型参数、优化器状态、梯度在 Flax 中几乎都是嵌套字典VariableDict / FrozenDict形态。然而实际工程中经常遇到两类需求整体拍平把任意深度的嵌套字典变成扁平字典键为路径元组方便按路径精确检索、排序、序列化或批量处理定向更新只对树中符合某种路径条件的子结构做变换且绝不原地修改——这是 JAX 函数式风格的核心约束。flax.traverse_util正是为解决这两类问题而存在。模块 docstring 开宗明义A utility for traversing immutable datastructures遍历不可变数据结构的工具并强调Traversals never mutate the original dataTraversal 永不修改原始数据一次update本质上是返回一份包含指定更新的数据副本。该模块位于 flax/traverse_util.py被 flax/training/checkpoints.py、flax/linen/module.py 等核心文件直接依赖是整个 Flax 基础设施的底层件之一。一、flatten_dict把嵌套字典拍平为路径键字典flatten_dict是模块中最常用的函数负责把嵌套字典展平。其 API 文档见 docs_nnx/api_reference/flax.traverse_util.rst 的Dict utils一节与 docstring 给出的核心示例如下from flax.traverse_util import flatten_dict xs {foo: 1, bar: {a: 2, b: {}}} flat_xs flatten_dict(xs) # flat_xs # {(foo,): 1, (bar, a): 2}注意空字典默认被忽略{bar: {b: {}}}中的空字典b不会出现在结果里unflatten_dict也无法还原它。参数详解源码签名flax/traverse_util.pydef flatten_dict(xs, keep_empty_nodesFalse, is_leafNone, sepNone):参数默认值作用xs必填输入的嵌套字典必须是dict或flax.core.FrozenDict否则触发assert断言失败源码第 137-139 行keep_empty_nodesFalse为True时空字典不再被丢弃而是以模块级哨兵对象traverse_util.empty_node作为值保留便于无损往返is_leafNone可选函数接收(prefix, xs)两个参数返回True表示当前嵌套字典应被视为叶子、不再继续展开sepNone若指定如/返回字典的键由路径元组改为sep连接而成的字符串None时键为元组keep_empty_nodes与empty_node哨兵源码中empty_node被定义为struct.dataclass class _EmptyNode: pass empty_node _EmptyNode()它特意用flax.struct.dataclass装饰注释明确说明是为了 be compatible with JAX与 JAX 兼容这样哨兵可以作为 pytree 叶子参与jax.jit等变换。开启keep_empty_nodesTrue后flat_xs flatten_dict(xs, keep_empty_nodesTrue) # {(foo,): 1, (bar, a): 2, (bar, b): traverse_util.empty_node} xs_restore unflatten_dict(flat_xs) # {foo: 1, bar: {a: 2, b: {}}} —— 与原始输入完全一致测试 tests/traverse_util_test.py 的test_flatten_dict_keep_empty验证了这一往返一致性。is_leaf自定义叶子判定当你想把某个深度的子字典整体当作叶子值保留时使用is_leaf。测试test_flatten_dict_is_leaftests/traverse_util_test.py展示了典型用法xs {foo: {c: 4}, bar: {a: 2, b: {}}} flat_xs flatten_dict( xs, is_leaflambda k, x: len(k) 1 and len(x) 2 ) # {(foo, c): 4, (bar,): {a: 2, b: {}}} # —— bar 因满足 len(路径)1 且 len(字典)2 而被视为叶子整体保留注意is_leaf的判定发生在递归展开之前当某节点既满足叶子条件又仍是 dict 时它整体作为一个值被保留其内部结构不再展开。sep路径字符串化测试test_flatten_dicttests/traverse_util_test.py验证了sep/的用法flat_xs flatten_dict(xs, sep/) # {foo: 1, bar/a: 2}这在需要把路径直接拼进文件名、日志标签或 GCS/tensorstore 子路径的场景中非常实用下文 checkpoint 部分会看到真实案例。二、unflatten_dict还原嵌套结构unflatten_dict是flatten_dict的逆操作flax/traverse_util.pyfrom flax.traverse_util import unflatten_dict flat_xs { (foo,): 1, (bar, a): 2, } xs unflatten_dict(flat_xs) # {foo: 1, bar: {a: 2}}要点输入必须是普通dict键为路径元组或字符串配合sep使用若值与empty_node相等还原时自动替换为{}对应keep_empty_nodesTrue的往返中途路径缺失时按需自动创建中间层字典与flatten_dict使用同一个sep参数路径字符串会按sep切分回元组。源码第 165-178 行展示了实现遍历每个(path, value)沿路径逐层下沉创建cursor字典最后把值挂到叶子键上。三、path_aware_map带路径信息的叶子级映射path_aware_map是路径感知的 mapflax/traverse_util.py它对嵌套字典的每个叶子调用f(path, value)其中path是该叶子从根到自身的路径元组。docstring 示例import jax.numpy as jnp from flax import traverse_util params {a: {x: 10, y: 3}, b: {x: 20}} f lambda path, x: x 5 if x in path else -x traverse_util.path_aware_map(f, params) # {a: {x: 15, y: -3}, b: {x: 25}}实现原理源码实现只有两行核心逻辑flat flatten_dict(nested_dict, keep_empty_nodesTrue) return unflatten_dict( {k: f(k, v) if v is not empty_node else v for k, v in flat.items()} )即先无损拍平keep_empty_nodesTrue保住空节点对每个叶子应用f再把结果还原为嵌套结构。因为flatten_dict/unflatten_dict都是纯函数式操作path_aware_map天然不修改输入并能在结果中保留空字典测试test_path_aware_map_with_empty_nodestests/traverse_util_test.py。实战与 optaxmulti_transform配合实现分层优化这是path_aware_map最典型的落地场景——为不同路径的参数贴上不同标签交给 optax 做多组优化器调度。测试test_path_aware_map_with_multi_transformtests/traverse_util_test.py给出了完整可运行示例params { linear_1: {w: jnp.zeros((5, 6)), b: jnp.zeros(5)}, linear_2: {w: jnp.zeros((6, 1)), b: jnp.zeros(1)}, } gradients jax.tree_util.tree_map(jnp.ones_like, params) # 占位梯度 # 按路径是否为 w 打标签kernel / bias param_labels traverse_util.path_aware_map( lambda path, x: kernel if w in path else bias, params ) tx optax.multi_transform( {kernel: optax.sgd(1.0), bias: optax.set_to_zero()}, param_labels ) state tx.init(params) updates, new_state tx.update(gradients, state, params) new_params optax.apply_updates(params, updates)效果w权重使用 SGD 更新b偏置完全冻结——最终b与原始值完全一致、w被更新。同样的模式也适用于optax.masked测试test_path_aware_map_with_maskedtests/traverse_util_test.py可见该函数是按参数名/路径做选择性优化的通用前置工具。四、Traversal 遍历 API可组合的不可变数据结构访问器模块 docstring 与源码共同定义了一套面向对象风格的 Traversal 体系。它的设计哲学是Traversal 是一个透镜lens——它选中数据结构中的一个子集支持iterate读取与update写入副本并可通过组合构建出任意复杂的选区。4.1 基础用法选中与更新从 identity 遍历t_identity出发通过方法链和属性/下标访问构造目标选区from flax import traverse_util import dataclasses dataclasses.dataclass class Foo: foo: int 0 bar: int 0 # 属性遍历选中 Foo 的 foo 属性 x Foo(foo1) iterator traverse_util.TraverseAttr(foo).iterate(x) list(iterator) # [1] # 组合遍历每个元素取其 foo 键 data [{foo: 1, bar: 2}, {foo: 3, bar: 4}] traversal traverse_util.t_identity.each()[foo] list(traversal.iterate(data)) # [1, 3] # update不修改原对象返回更新后的副本 data {foo: Foo(bar2)} traversal traverse_util.t_identity[foo].bar data traversal.update(lambda x: x x, data) # {foo: Foo(foo0, bar4)}4.2 Traversal 类族与组合原语核心抽象类Traversalflax/traverse_util.py定义了两个抽象方法并提供了组合方法方法作用update(fn, inputs)对选中的每个元素应用fn返回新对象不可变风格iterate(inputs)返回一个迭代器产出选中的元素set(values, inputs)用一组新值覆盖选区值数量不匹配时抛ValueErrorcompose(other)组合两个 Traversal先走外层、再走内层merge(*traversals)合并多个 Traversal 的选区用于同时选中多个分支each()遍历容器中的每个元素dict/list/tupletree()遍历 pytree 的每个叶子filter(fn)按谓词过滤选中的值__getattr__/__getitem__语法糖.attr等价于compose(TraverseAttr(attr))[key]等价于compose(TraverseItem(key))具体实现类包括TraverseId恒等遍历iterate产出自身update直接应用fn全局单例t_identity是它源码第 296-308 行TraverseAttr遍历对象的属性支持 namedtuple_replace、dataclassdataclasses.replace及普通对象copy.copy后setattrTraverseItem遍历容器下标或键支持元组、namedtuple、list、dict 以及slice 切片t_identity[1:3]可一次选中多个元组元素TraverseEach遍历 list/tuple/dict 中每个条目对非这三种类型抛ValueErrorTraverseTree基于jax.tree_util遍历 pytree 所有叶子TraverseFilter按谓词过滤TraverseMerge将多个 Traversal 的选区合并成一个选区TraverseCompose串联两个 Traversal是上述所有组合方法的底层引擎。4.3 测试覆盖的典型组合tests/traverse_util_test.py 的TraversalTest类逐条验证了这些行为# 元组下标与切片 x (1, 2, 3, 4) list(traverse_util.t_identity[1:3].iterate(x)) # [2, 3] traverse_util.t_identity[1:3].update(lambda x: x x, x) # (1, 4, 6, 4) # namedtuple 属性 Point collections.namedtuple(Point, [x, y]) x Point(x1, y2) traverse_util.t_identity.y.update(lambda x: x x, x) # Point(x1, y4) # each merge同时处理多个字段 x [{foo: 1, bar: 2}, {foo: 3, bar: 4}] t traverse_util.t_identity.each().merge( traverse_util.TraverseItem(foo), traverse_util.TraverseItem(bar) ) list(t.iterate(x)) # [1, 2, 3, 4] # filter按条件过滤后更新 x [1, -2, 3, -4] t traverse_util.t_identity.each().filter(lambda x: x 0) t.update(lambda x: -x, x) # [1, 2, 3, 4] # set校验值数量 t traverse_util.t_identity[foo].each() t.set([3, 4], {foo: [1, 2]}) # {foo: [3, 4]} # 值太少 / 太多都会抛 ValueError4.4 弃用提示Traversal 与flax.optim源码第 214-223 行在Traversal.__new__中发出DeprecationWarning提示flax.traverse_util.Traversal将被弃用。如果你是为了flax.optim使用它请改用optax详细迁移说明见官方 optax 更新指南docs/guides/converting_and_upgrading/optax_update_guide.rst 可找到对应的迁移思路。也就是说新代码中处理优化器相关需求请优先使用 optax 及上文介绍的path_aware_map方案Traversal 类目前仍保留t_identity定义时用warnings.catch_warnings()抑制了弃用告警保证其可用主要服务于旧的flax.optim兼容路径与更广义的不可变结构遍历需求。4.5ModelParamTraversal按参数全名筛选ModelParamTraversalflax/traverse_util.py专为模型参数设计构造时传入filter_fn它会收到形如/module/sub_module/parameter_name的完整路径字符串与参数值返回该参数是否被选中。测试test_param_selectiontests/traverse_util_test.py展示了按名称包含kernel筛选并翻倍更新的效果params {x: {kernel: 1, bias: 2, y: {kernel: 3, bias: 4}, z: {}}} traversal traverse_util.ModelParamTraversal( lambda name, _: kernel in name ) list(traversal.iterate(params)) # [1, 3]只选中两个 kernel traversal.update(lambda x: x x, params) # kernel 变为 2、6bias 不变它内部通过flatten_dict含keep_empty_nodesTrue获取扁平参数表按字典序排序后处理对flax.core.FrozenDict输入会返回FrozenDict结果保持不可变类型不变。同时它只接受嵌套 dict 或FrozenDict对其他类型抛ValueError测试test_only_works_on_model_params验证了这一点。五、仓库内的真实应用checkpoint 与 Linen 内部的依赖traverse_util并非孤立工具它在 Flax 基础设施中承担关键角色以下用法均有源码可查5.1 checkpoint 中的多进程数组MPA分片处理flax/training/checkpoints.py 的_split_mp_arrays在保存 checkpoint 时用flatten_dict(target, keep_empty_nodesTrue)拍平整个目标树挑出所有多进程分布式数组把路径/.join(key)拼成子路径并替换为占位符最后用unflatten_dict还原flattened traverse_util.flatten_dict(target, keep_empty_nodesTrue) mpa_targets [] for key, value in flattened.items(): if _is_multiprocess_array(value): subpath /.join(key) mpa_targets.append((value, subpath)) flattened[key] MP_ARRAY_PH subpath target traverse_util.unflatten_dict(flattened)恢复路径_restore_mpasflax/training/checkpoints.py同样依赖flatten_dict/unflatten_dict完成占位符替换与还原。这正是sep参数与keep_empty_nodes在真实工程中的价值体现。5.2 Linen 模块中对叶子或整个 dict 统一映射flax/linen/module.py 在map_相关逻辑中用flatten_dict普通版与keep_empty_nodesTrue版处理叶子值与嵌套字典混合的输入再以unflatten_dict恢复结构。flax.core.lift的 remat/scan 内部flax/core/lift.py也用它做中间表示的拍平与还原。这些使用点证明了该模块是 Flax 参数树序列化与变换管线的标准工具理解它能帮助你读懂 checkpoint 源码乃至自定义序列化逻辑。六、快速参考API 一览与选择建议API一句话定位典型场景flatten_dict(xs, keep_empty_nodes, is_leaf, sep)嵌套字典 → 路径键元组或字符串字典参数树拍平、按路径排序/检索、路径拼接unflatten_dict(xs, sepNone)路径键字典 → 嵌套字典还原拍平结果、checkpoint 占位符回填path_aware_map(f, nested_dict)对每个叶子调用f(path, value)并还原参数打标签喂给optax.multi_transform/optax.masked、按路径做条件变换t_identity/ Traversal 族可组合的不可变选区透镜选中并更新嵌套结构中的子集注意弃用提示ModelParamTraversal(filter_fn)按/module/param全名筛选参数传统flax.optim多优化器场景的参数分组选型建议新项目处理按路径映射优先用path_aware_map纯函数、无弃用风险、与 optax 天然契合需要整体扁平化做序列化/哈希/排序时用flatten_dictunflatten_dict注意空字典与keep_empty_nodes的取舍Traversal 类族适用于需要精确读写对象属性/切片/多分支合并的不可变结构操作但请知晓其DeprecationWarning背景。七、易错点与工程注意事项flatten_dict只接受 dict / FrozenDict传入 list、tuple 或其他容器会直接触发assert失败源码 137-139 行这类结构请先用jax.tree_util或在外面包一层 dict。空字典默认丢失若你的结构中有空 dict 且需要无损往返务必设置keep_empty_nodesTrue否则unflatten_dict无法恢复原状。unflatten_dict输入必须是普通 dict传入 FrozenDict 会触发断言如需保持 FrozenDict 类型可参考ModelParamTraversal的处理方式返回时手动包回FrozenDict。path_aware_map中f的签名是(path, value)与jax.tree_util.tree_map的单参数f(value)不同path是字符串元组可用于判断路径是否包含某个字段名。Traversal 的不可变性update返回新对象而不改原对象但需要注意TraverseItem对 list/dict 使用copy.copy的浅拷贝语义——叶子元素被替换但未被选中的嵌套子结构仍与原对象共享引用。版本背景本文描述的行为以当前仓库 flax/traverse_util.pyCopyright 2024为准Traversal 相关 API 处于弃用过渡期新代码应避开依赖弃用告警的调用方式。通过上述六个章节你已掌握flax.traverse_util的全部核心 API、底层实现机理、测试验证路径以及在 checkpoint 与 Linen 内部的真实工程用法——这套工具虽然轻量却是理解 Flax 参数树操作流水线的重要基石。【免费下载链接】flaxFlax is a neural network library for JAX that is designed for flexibility.项目地址: https://gitcode.com/GitHub_Trending/fl/flax创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考