迁移到普通字典

当调用 .init、.init_with_output 和 .apply Module 方法时，Flax 将从返回 FrozenDicts 迁移到返回普通字典。

原始问题概述于 此处。

本指南展示了一些常见的升级模式。

实用程序函数

FrozenDicts 是不可变的字典，它们实现了另外 4 个方法

• copy

• pop

• pretty_repr

• unfreeze

为了适应普通字典的变化，请将 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（在此处找到 here）。

迁移状态

截至 2023 年 7 月 19 日，flax_return_frozendict 设置为 False（见 #3193），这意味着从版本 0.7.1 开始，Flax 将默认从返回普通字典。此标志可以暂时翻转为 True 以使 Flax 返回 Frozendicts。但是，此特性标志最终将在未来被删除。