移轉到常规的詞典

當呼叫 .init、 .init_with_output 和 .apply Module 方法時，Flax 將從回傳 FrozenDicts 遷移為一般字典。

原始議題在此 處說明。

本指南顯示一些常見的升級模式。

工具函式

FrozenDicts 是不可變字典，此字典實作了其他 4 個方法

• 複製

• 彈出

• pretty_repr

• 解凍

為了配合一般字典的變更，請將 FrozenDict 方法的使用替換成 flax.core.frozen_dict 的工具函式等效項。這些工具函式模仿其對應 FrozenDict 方法的行為，且可以在 FrozenDicts 或一般字典中呼叫。以下是工具函式和範例升級模式

copy

僅限於 FrozenDict

variables = variables.copy(add_or_replace={'other_variables': other_variables})

FrozenDict 和一般字典皆可

variables = flax.core.copy(variables, add_or_replace={'other_variables': other_variables})

pop

僅限於 FrozenDict

state, params = variables.pop('params')

FrozenDict 和一般字典皆可

state, params = flax.core.pop(variables, 'params')

pretty_repr

僅限於 FrozenDict

str_repr = variables.pretty_repr()

FrozenDict 和一般字典皆可

str_repr = flax.core.pretty_repr(variables)

unfreeze

僅限於 FrozenDict

variables = variables.unfreeze()

FrozenDict 和一般字典皆可

variables = flax.core.unfreeze(variables)

修改設定值

設定暫時功能旗標 flax_return_frozendict 以協助執行遷移。如要在執行期間切換傳回 FrozenDict 與常規字典變數的行為，請在程式碼中執行 flax.config.update('flax_return_frozendict', <BOOLEAN_VALUE>)。

例如

x = jnp.empty((1,3))

flax.config.update('flax_return_frozendict', True) # set Flax to return FrozenDicts
variables = nn.Dense(5).init(jax.random.key(0), x)

assert isinstance(variables, flax.core.FrozenDict)

flax.config.update('flax_return_frozendict', False) # set Flax to return regular dicts
variables = nn.Dense(5).init(jax.random.key(0), x)

assert isinstance(variables, dict)

或者，可直接修改 Flax 原始碼中的環境變數 flax_return_frozendict（請參閱此處）

遷移狀態

截至 2023 年 7 月 19 日，flax_return_frozendict 已設為 False（請參閱 #3193），表示 Flax 將預設傳回 v0.7.1 之後的版本中的常規字典。此旗標可以臨時切換為 True 以讓 Flax 傳回 Frozendicts。不過，此功能旗標最終將在未來移除。