装饰器

装饰器#

flax.linen.compact(fun)[源代码]#

标记给定的模块方法，允许内联子模块。

用 @compact 包装的方法可以直接在方法内定义子模块。

例如

>>> import flax.linen as nn

>>> class Foo(nn.Module):
...   @nn.compact
...   def __call__(self, x, features):
...     x = nn.Dense(features)(x)
...     ...
...     return x

每个模块中最多只能有一个方法用 @compact 包装。

参数

fun – 要标记为 compact 的模块方法。

返回值

给定的函数 fun 被标记为 compact。

flax.linen.nowrap(fun)[源代码]#

将给定的模块方法标记为不需要包装的辅助方法。

用 @nowrap 包装的方法是私有辅助方法，不需要用状态处理程序或单独的 named_call 转换进行包装。

这在几个具体实例中是必要的

  • 如果您正在子类化像 Module.param 这样的方法，并且不希望这个被覆盖的核心函数用状态管理包装器装饰。

  • 如果您希望一个方法可以从一个未绑定的模块调用（例如：不依赖于参数/RNG 的参数构造函数）。如果您想了解更多关于 Flax 模块如何管理其状态的信息，请阅读 [Flax 模块生命周期](https://flax.org.cn/en/latest/developer_notes/module_lifecycle.html) 指南。

例如

>>> import flax.linen as nn
>>> import jax, jax.numpy as jnp

>>> class Foo(nn.Module):
...   num_features: int

...   @nn.nowrap
...   def _make_dense(self, num_features):
...     return nn.Dense(num_features)

...   @nn.compact
...   def __call__(self, x):
...     # now safe to use constructor helper even if using named_call
...     dense = self._make_dense(self.num_features)
...     return dense(x)
参数

fun – 要标记为 nowrap 的模块方法。

返回值

给定的函数 fun 被标记为 nowrap。

目录