2026/9/16 22:52:18
FEATURED · 精选文章 
Flax Linen Init/Apply API深入解析 flax.linen.apply、init 与 init_with_output【免费下载链接】flaxFlax is a neural network library for JAX that is designed for flexibility.项目地址: https://gitcode.com/GitHub_Trending/fl/flaxFlax 的flax.linen除了提供Module类方法式的module.init(...)/module.apply(...)之外还暴露了一组函数式入口模块级函数apply、init和init_with_output。本文基于 Init/Apply 官方 API 文档 及其在 flax/linen/module.py 中的真实实现系统讲解这三个函数的签名、参数语义、与Module方法的区别以及它们如何落到底层flax.core的 Scope 机制上帮助你理解 Flax “初始化/推理分离”这一核心编程模型的完整调用链。一、Init/ApplyFlax 编程模型的核心页Init/Apply 是 Flax Linen API 参考中专门记录这一组函数的页面其内容由三个自动生成的函数文档组成.. autofunction:: apply .. autofunction:: init .. autofunction:: init_with_output这组函数的存在意义是把 Flax 的“绑定bound/未绑定unbound”两种使用方式衔接起来Module.init/Module.apply是绑定式的——你必须先持有模块实例再调用方法而flax.linen.apply等函数式入口可以先编写一个普通的纯函数例如组合多个模块的 encode/decode 流程再把它“函数化”成可接受外部variables的 callable从而直接参与jax.jit、jax.grad等 JAX 变换。二、三个函数的签名速览三个函数均在 flax/linen/module.py 中定义源码签名如下# 源码位置: flax/linen/module.py def apply( fn: Callable[..., Any], module: Module, mutable: CollectionFilter False, capture_intermediates: bool | Callable[[Module, str], bool] False, ) - Callable[..., Any]: ... def init_with_output( fn: Callable[..., Any], module: Module, mutable: CollectionFilter DenyList(intermediates), capture_intermediates: bool | Callable[[Module, str], bool] False, ) - Callable[..., tuple[Any, FrozenVariableDict | dict[str, Any]]]: ... def init( fn: Callable[..., Any], module: Module, mutable: CollectionFilter DenyList(intermediates), capture_intermediates: bool | Callable[[Module, str], bool] False, ) - Callable[..., FrozenVariableDict | dict[str, Any]]: ...可以归纳为一张对照表函数返回的 callable 签名返回类型mutable默认值flax.linen.apply(variables, *args, rngsNone, **kwargs) - TT若mutable非False则为(T, 变更后的变量)元组Falseflax.linen.init_with_output(rngs, *args, **kwargs) - (T, variables)输出与变量的元组DenyList(intermediates)flax.linen.init(rngs, *args, **kwargs) - variables仅变量DenyList(intermediates)其中T是被包装函数fn的返回类型variables是按集合collection组织的变量字典即文档中提到的FrozenVariableDict由 flax/core/frozen_dict.py 中的FrozenDict提供也可以是与普通dict兼容的结构。rngs既可以是PRNGKey字典也可以是单个PRNGKey——后者等价于传入名为params的单个键。三、flax.linen.apply把 fn 变成接收 variables 的函数apply的核心约定见 apply 源码 的 docstring传入一个普通函数fn其第一个形参将被视为与module相同的模块实例实际是一个 clone下文详述返回一个新函数签名为(variables, *args, rngsNone, **kwargs) - T若mutable不是False返回类型变为二元组第二项是FrozenDict形式的变更后的变量。官方 docstring 给出的标准示例可直接复制运行class Foo(nn.Module): def encode(self, x): ... def decode(self, x): ... def f(foo, x): z foo.encode(x) y foo.decode(z) # ... return y variables {} foo Foo() f_jitted jax.jit(nn.apply(f, foo)) f_jitted(variables, jnp.ones((1, 3)))注意最后一步jax.jit直接包裹函数式入口的返回值这正是这类 API 与“先构造模块再 apply”写法的关键差异——fn中可以对模块自由编排先encode再decode甚至穿插其他 JAX 变换而整个过程仍是纯函数、可追踪的。参数语义fn要执行的函数。传入它的第一实参是“绑定了 variables 与 RNG 的模块实例”且是module的克隆源码实现中通过module.clone(parentscope, _deep_cloneTrue)完成module用于绑定变量与 RNG 的原型模块mutable可取bool、str或集合名的list指定哪些 collection 可写。bool表示全部/不允许可变str表示单个集合名list表示多个集合名capture_intermediates为True时把各子模块__call__的中间输出捕获到intermediates集合中也可传入一个(module, method_name) - bool的过滤函数精细控制哪些方法的输出被捕获。四、flax.linen.init 与 init_with_output初始化侧的函数式入口init_with_output源码与apply对称返回的 callable 签名是(rngs, *args, **kwargs) - (T, variables)即同时给出函数输出和初始化出来的变量。这在需要初始化时顺便取一次前向输出例如检查 loss 初始值、生成第一个 batch 的预测时很有用foo Foo() f_jitted jax.jit(nn.init_with_output(f, foo)) y, variables f_jitted(jax.random.key(0), jnp.ones((1, 3)))init源码则是init_with_output的薄封装丢弃输出、只返回变量 foo Foo() f_jitted jax.jit(nn.init(f, foo)) variables f_jitted(jax.random.key(0), jnp.ones((1, 3)))从源码结构看init的实现完全委托给init_with_output先构造init_fn再用一个init_wrapper取init_fn(*args, **kwargs)[1]。这也解释了两者共享相同的mutable与capture_intermediates参数以及相同的默认值DenyList(intermediates)——初始化时允许写入所有集合唯独排除intermediates该集合是capture_intermediates机制的专属命名空间不应作为普通参数/状态持久化。五、与 Module 方法式 API 的关系Linen 中每个函数都有一个对应的Module实例方法Module.apply源码、Module.init_with_output源码 起、Module.init源码 起。docstring 明确写道 “UnlikeModule.applythis function returns a new function...”即二者的区别在于抽象层级Module.apply(variables, *args, rngs..., method..., mutable...)直接对模块实例的__call__或method指定的方法做函数化适合“一个模块 一个模型”的常规场景flax.linen.apply(fn, module)以任意纯函数为组合单元适合把多个模块、采样逻辑、后处理串成一个fn再整体函数化。值得一提的是Module.apply本身在实现末尾也是调用函数式入口它做若干预处理rngs单键归一化为{params: rngs}、字符串方法名解析为可调用对象等后执行apply(method, self, mutablemutable, capture_intermediates...)(variables, *args, **kwargs, rngsrngs)见 Module.apply 尾部。因此可以说函数式入口才是真正干活的层方法式 API 是它的便捷封装。Module.apply的 docstring 还补充了几条函数式入口同样适用的运行时语义值得注意RNG 流回退若传入单个PRNGKeyFlax 用它供给params流self.make_rng(name)请求的流名若未在rngs中给出将回退使用params流。docstring 中的示例验证了这一点rngs里删掉noise键后make_rng(noise)的行为与直接传单个 key 一致method 的多态方法式 API 的method参数可以是不绑定方法、字符串甚至模块外定义的函数该函数需以模块实例为首参——函数式入口的fn本质上就是最后这种“模块外函数”的自由形式。六、底层调用链从 linen 到 core.Scope三个函数式入口内部结构高度一致以apply为例见 源码functools.wraps(fn) def scope_fn(scope, *args, **kwargs): _context.capture_stack.append(capture_intermediates) try: return fn(module.clone(parentscope, _deep_cloneTrue), *args, **kwargs) finally: _context.capture_stack.pop() if capture_intermediates is True: capture_intermediates capture_call_intermediates if capture_intermediates: mutable union_filters(mutable, intermediates) return core.apply(scope_fn, mutablemutable)可以归纳出四条实现事实函数化的对象是scope_fn而非fn。scope_fn接收一个flax.core的Scope把module深克隆并挂到该 scope 上parentscope, _deep_cloneTrue再转发调用fn最终委托给flax.core。core.apply(scope_fn, mutablemutable)/core.init(scope_fn, mutablemutable)指向 flax/core/scope.py 中的flax.core.apply/flax.core.init其 docstring 描述为 “Functionalize aScopefunction”——即 Linen 层把“绑定 Module”的问题规约成 core 层“绑定 Scope”的问题中间量捕获是线程本地状态capture_intermediates通过_context.capture_stack压栈/弹栈传递并在启用时自动用union_filters把intermediates并入mutable过滤器保证捕获目标可写capture_intermediatesTrue的语义是“只捕获__call__”True会被替换为内置的capture_call_intermediates过滤器若要捕获非__call__方法的输出需传入自定义的(module, method_name) - bool过滤函数。七、实战示例函数式入口 jit 的完整闭环综合上述语义一个典型的函数式训练片段encode/decode结构即 docstring 原例import jax import jax.numpy as jnp import flax.linen as nn class Encoder(nn.Module): nn.compact def __call__(self, x): return nn.Dense(8)(x) class Decoder(nn.Module): nn.compact def __call__(self, z): return nn.Dense(1)(z) def loss_fn(model, x): z model.encode(x) y model.decode(z) return jnp.mean((y - x) ** 2) model nn.Sequential(Encoder(), Decoder()) # 任意 nn.Module 原型 # 初始化只取变量 init_jit jax.jit(nn.init(loss_fn, model)) variables init_jit(jax.random.key(0), jnp.ones((1, 1))) # 推理/前向直接消费 variables apply_jit jax.jit(nn.apply(loss_fn, model)) loss apply_jit(variables, jnp.ones((1, 1)))要点回顾nn.init(loss_fn, model)生成的 callable 只认(rngs, *args)不要求事先存在variables而nn.apply(loss_fn, model)生成的 callable 只认(variables, *args)初始化与推理在数据流上完全解耦——这正是 Flax 能够自由组合jit/grad、并对变量做序列化FrozenDict可直接持久化的根本原因。八、参考位置汇总API 文档页docs/api_reference/flax.linen/init_apply.rst其索引见 docs/api_reference/flax.linen/index.rst函数式入口实现flax/linen/module.py 中的applyL2968、init_with_outputL3038、initL3109方法式对应实现同文件中的Module.applyL2092、Module.init_with_outputL2252、Module.initL2316底层 Scope 函数化flax/core/scope.py 中的flax.core.applyL1050与flax.core.initL1103变量容器flax/core/frozen_dict.pyFrozenDict、freeze、unfreeze相关测试nn.init(的用法可在 tests/linen/linen_module_test.py 与 tests/linen/linen_recurrent_test.py 中检索到可作为回归行为的对照文档构建配置docs/conf.pyAPI 参考由 Sphinxautofunction自动从源码 docstring 生成。需要说明的适用前提本文所述行为以当前仓库的flax.linen源码为准。init/init_with_output的mutable默认值DenyList(intermediates)与apply的默认值False不同调用时若自定义了 collection应显式传入mutable过滤条件避免依赖默认语义。【免费下载链接】flaxFlax is a neural network library for JAX that is designed for flexibility.项目地址: https://gitcode.com/GitHub_Trending/fl/flax创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
Flax Linen Init/Apply API：深入解析 flax.linen.apply、init 与 init_with_output
Flax Linen Init/Apply API深入解析 flax.linen.apply、init 与 init_with_output【免费下载链接】flaxFlax is a neural network library for JAX that is designed for flexibility.项目地址: https://gitcode.com/GitHub_Trending/fl/flaxFlax 的flax.linen除了提供Module类方法式的module.init(...)/module.apply(...)之外还暴露了一组函数式入口模块级函数apply、init和init_with_output。本文基于 Init/Apply 官方 API 文档 及其在 flax/linen/module.py 中的真实实现系统讲解这三个函数的签名、参数语义、与Module方法的区别以及它们如何落到底层flax.core的 Scope 机制上帮助你理解 Flax “初始化/推理分离”这一核心编程模型的完整调用链。一、Init/ApplyFlax 编程模型的核心页Init/Apply 是 Flax Linen API 参考中专门记录这一组函数的页面其内容由三个自动生成的函数文档组成.. autofunction:: apply .. autofunction:: init .. autofunction:: init_with_output这组函数的存在意义是把 Flax 的“绑定bound/未绑定unbound”两种使用方式衔接起来Module.init/Module.apply是绑定式的——你必须先持有模块实例再调用方法而flax.linen.apply等函数式入口可以先编写一个普通的纯函数例如组合多个模块的 encode/decode 流程再把它“函数化”成可接受外部variables的 callable从而直接参与jax.jit、jax.grad等 JAX 变换。二、三个函数的签名速览三个函数均在 flax/linen/module.py 中定义源码签名如下# 源码位置: flax/linen/module.py def apply( fn: Callable[..., Any], module: Module, mutable: CollectionFilter False, capture_intermediates: bool | Callable[[Module, str], bool] False, ) - Callable[..., Any]: ... def init_with_output( fn: Callable[..., Any], module: Module, mutable: CollectionFilter DenyList(intermediates), capture_intermediates: bool | Callable[[Module, str], bool] False, ) - Callable[..., tuple[Any, FrozenVariableDict | dict[str, Any]]]: ... def init( fn: Callable[..., Any], module: Module, mutable: CollectionFilter DenyList(intermediates), capture_intermediates: bool | Callable[[Module, str], bool] False, ) - Callable[..., FrozenVariableDict | dict[str, Any]]: ...可以归纳为一张对照表函数返回的 callable 签名返回类型mutable默认值flax.linen.apply(variables, *args, rngsNone, **kwargs) - TT若mutable非False则为(T, 变更后的变量)元组Falseflax.linen.init_with_output(rngs, *args, **kwargs) - (T, variables)输出与变量的元组DenyList(intermediates)flax.linen.init(rngs, *args, **kwargs) - variables仅变量DenyList(intermediates)其中T是被包装函数fn的返回类型variables是按集合collection组织的变量字典即文档中提到的FrozenVariableDict由 flax/core/frozen_dict.py 中的FrozenDict提供也可以是与普通dict兼容的结构。rngs既可以是PRNGKey字典也可以是单个PRNGKey——后者等价于传入名为params的单个键。三、flax.linen.apply把 fn 变成接收 variables 的函数apply的核心约定见 apply 源码 的 docstring传入一个普通函数fn其第一个形参将被视为与module相同的模块实例实际是一个 clone下文详述返回一个新函数签名为(variables, *args, rngsNone, **kwargs) - T若mutable不是False返回类型变为二元组第二项是FrozenDict形式的变更后的变量。官方 docstring 给出的标准示例可直接复制运行class Foo(nn.Module): def encode(self, x): ... def decode(self, x): ... def f(foo, x): z foo.encode(x) y foo.decode(z) # ... return y variables {} foo Foo() f_jitted jax.jit(nn.apply(f, foo)) f_jitted(variables, jnp.ones((1, 3)))注意最后一步jax.jit直接包裹函数式入口的返回值这正是这类 API 与“先构造模块再 apply”写法的关键差异——fn中可以对模块自由编排先encode再decode甚至穿插其他 JAX 变换而整个过程仍是纯函数、可追踪的。参数语义fn要执行的函数。传入它的第一实参是“绑定了 variables 与 RNG 的模块实例”且是module的克隆源码实现中通过module.clone(parentscope, _deep_cloneTrue)完成module用于绑定变量与 RNG 的原型模块mutable可取bool、str或集合名的list指定哪些 collection 可写。bool表示全部/不允许可变str表示单个集合名list表示多个集合名capture_intermediates为True时把各子模块__call__的中间输出捕获到intermediates集合中也可传入一个(module, method_name) - bool的过滤函数精细控制哪些方法的输出被捕获。四、flax.linen.init 与 init_with_output初始化侧的函数式入口init_with_output源码与apply对称返回的 callable 签名是(rngs, *args, **kwargs) - (T, variables)即同时给出函数输出和初始化出来的变量。这在需要初始化时顺便取一次前向输出例如检查 loss 初始值、生成第一个 batch 的预测时很有用foo Foo() f_jitted jax.jit(nn.init_with_output(f, foo)) y, variables f_jitted(jax.random.key(0), jnp.ones((1, 3)))init源码则是init_with_output的薄封装丢弃输出、只返回变量 foo Foo() f_jitted jax.jit(nn.init(f, foo)) variables f_jitted(jax.random.key(0), jnp.ones((1, 3)))从源码结构看init的实现完全委托给init_with_output先构造init_fn再用一个init_wrapper取init_fn(*args, **kwargs)[1]。这也解释了两者共享相同的mutable与capture_intermediates参数以及相同的默认值DenyList(intermediates)——初始化时允许写入所有集合唯独排除intermediates该集合是capture_intermediates机制的专属命名空间不应作为普通参数/状态持久化。五、与 Module 方法式 API 的关系Linen 中每个函数都有一个对应的Module实例方法Module.apply源码、Module.init_with_output源码 起、Module.init源码 起。docstring 明确写道 “UnlikeModule.applythis function returns a new function...”即二者的区别在于抽象层级Module.apply(variables, *args, rngs..., method..., mutable...)直接对模块实例的__call__或method指定的方法做函数化适合“一个模块 一个模型”的常规场景flax.linen.apply(fn, module)以任意纯函数为组合单元适合把多个模块、采样逻辑、后处理串成一个fn再整体函数化。值得一提的是Module.apply本身在实现末尾也是调用函数式入口它做若干预处理rngs单键归一化为{params: rngs}、字符串方法名解析为可调用对象等后执行apply(method, self, mutablemutable, capture_intermediates...)(variables, *args, **kwargs, rngsrngs)见 Module.apply 尾部。因此可以说函数式入口才是真正干活的层方法式 API 是它的便捷封装。Module.apply的 docstring 还补充了几条函数式入口同样适用的运行时语义值得注意RNG 流回退若传入单个PRNGKeyFlax 用它供给params流self.make_rng(name)请求的流名若未在rngs中给出将回退使用params流。docstring 中的示例验证了这一点rngs里删掉noise键后make_rng(noise)的行为与直接传单个 key 一致method 的多态方法式 API 的method参数可以是不绑定方法、字符串甚至模块外定义的函数该函数需以模块实例为首参——函数式入口的fn本质上就是最后这种“模块外函数”的自由形式。六、底层调用链从 linen 到 core.Scope三个函数式入口内部结构高度一致以apply为例见 源码functools.wraps(fn) def scope_fn(scope, *args, **kwargs): _context.capture_stack.append(capture_intermediates) try: return fn(module.clone(parentscope, _deep_cloneTrue), *args, **kwargs) finally: _context.capture_stack.pop() if capture_intermediates is True: capture_intermediates capture_call_intermediates if capture_intermediates: mutable union_filters(mutable, intermediates) return core.apply(scope_fn, mutablemutable)可以归纳出四条实现事实函数化的对象是scope_fn而非fn。scope_fn接收一个flax.core的Scope把module深克隆并挂到该 scope 上parentscope, _deep_cloneTrue再转发调用fn最终委托给flax.core。core.apply(scope_fn, mutablemutable)/core.init(scope_fn, mutablemutable)指向 flax/core/scope.py 中的flax.core.apply/flax.core.init其 docstring 描述为 “Functionalize aScopefunction”——即 Linen 层把“绑定 Module”的问题规约成 core 层“绑定 Scope”的问题中间量捕获是线程本地状态capture_intermediates通过_context.capture_stack压栈/弹栈传递并在启用时自动用union_filters把intermediates并入mutable过滤器保证捕获目标可写capture_intermediatesTrue的语义是“只捕获__call__”True会被替换为内置的capture_call_intermediates过滤器若要捕获非__call__方法的输出需传入自定义的(module, method_name) - bool过滤函数。七、实战示例函数式入口 jit 的完整闭环综合上述语义一个典型的函数式训练片段encode/decode结构即 docstring 原例import jax import jax.numpy as jnp import flax.linen as nn class Encoder(nn.Module): nn.compact def __call__(self, x): return nn.Dense(8)(x) class Decoder(nn.Module): nn.compact def __call__(self, z): return nn.Dense(1)(z) def loss_fn(model, x): z model.encode(x) y model.decode(z) return jnp.mean((y - x) ** 2) model nn.Sequential(Encoder(), Decoder()) # 任意 nn.Module 原型 # 初始化只取变量 init_jit jax.jit(nn.init(loss_fn, model)) variables init_jit(jax.random.key(0), jnp.ones((1, 1))) # 推理/前向直接消费 variables apply_jit jax.jit(nn.apply(loss_fn, model)) loss apply_jit(variables, jnp.ones((1, 1)))要点回顾nn.init(loss_fn, model)生成的 callable 只认(rngs, *args)不要求事先存在variables而nn.apply(loss_fn, model)生成的 callable 只认(variables, *args)初始化与推理在数据流上完全解耦——这正是 Flax 能够自由组合jit/grad、并对变量做序列化FrozenDict可直接持久化的根本原因。八、参考位置汇总API 文档页docs/api_reference/flax.linen/init_apply.rst其索引见 docs/api_reference/flax.linen/index.rst函数式入口实现flax/linen/module.py 中的applyL2968、init_with_outputL3038、initL3109方法式对应实现同文件中的Module.applyL2092、Module.init_with_outputL2252、Module.initL2316底层 Scope 函数化flax/core/scope.py 中的flax.core.applyL1050与flax.core.initL1103变量容器flax/core/frozen_dict.pyFrozenDict、freeze、unfreeze相关测试nn.init(的用法可在 tests/linen/linen_module_test.py 与 tests/linen/linen_recurrent_test.py 中检索到可作为回归行为的对照文档构建配置docs/conf.pyAPI 参考由 Sphinxautofunction自动从源码 docstring 生成。需要说明的适用前提本文所述行为以当前仓库的flax.linen源码为准。init/init_with_output的mutable默认值DenyList(intermediates)与apply的默认值False不同调用时若自定义了 collection应显式传入mutable过滤条件避免依赖默认语义。【免费下载链接】flaxFlax is a neural network library for JAX that is designed for flexibility.项目地址: https://gitcode.com/GitHub_Trending/fl/flax创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
