裝飾器

裝飾器#

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 – 將 Module 方法標記為精簡。

回傳值

指定函式 fun 標記為精簡。

flax.linen.nowrap(fun)[原始碼]#

將指定的模組方法標記為不需要封裝的輔助方法。

包裝在 @nowrap 中的方法是私人輔助方法，不需要用狀態處理程式或另外的 named_call 轉換封裝。

在下列具體情況中需要這樣做

  • 如果您要建立像 Module.param 之類的方法子類別，而且不想用狀態管理封裝程式來裝飾這個覆寫的核心函式。

  • 如果您想要一個方法可以從未繫結的 Module 呼叫（例如：建構參數時使用的函式，與參數/RNG 無關）。如果您想要深入瞭解 Flax Module 如何管理它們的狀態，請參閱 [The Flax Module lifecycle](https://flax.dev.org.tw/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 – 將 Module 方法標記為 nowrap。

回傳值

指定函式 fun 標記為 nowrap。

內容