装饰器

装饰器#

flax.linen.compact(fun)[source]#

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

用 @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 – 要标记为紧凑的模块方法。

返回值

标记为紧凑的给定函数 fun。

flax.linen.nowrap(fun)[source]#

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

用 @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 的模块方法。

返回值

标记为 nowrap 的给定函数 fun。

内容