flax.core.frozen_dict package

class flax.core.frozen_dict.FrozenDict(*args, __unsafe_skip_copy__=False, **kwargs)

An immutable variant of the Python dict.

copy(add_or_replace)

Create a new FrozenDict with additional or replaced entries.

Parameters

add_or_replace (Mapping[flax.core.frozen_dict.K, flax.core.frozen_dict.V]) –

Return type

flax.core.frozen_dict.FrozenDict[flax.core.frozen_dict.K, flax.core.frozen_dict.V]

pop(key)

Create a new FrozenDict where one entry is removed.

Example:

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

Parameters

key (flax.core.frozen_dict.K) – the key to remove from the dict

Returns

A pair with the new FrozenDict and the removed value.

Return type

Tuple[flax.core.frozen_dict.FrozenDict[flax.core.frozen_dict.K, flax.core.frozen_dict.V], flax.core.frozen_dict.V]

pretty_repr(num_spaces=4)

Returns an indented representation of the nested dictionary.

tree_flatten()

Flattens this FrozenDict.

Returns

A flattened version of this FrozenDict instance.

Return type

Tuple[Tuple[Dict[Any, Any]], Tuple[()]]

unfreeze()

Unfreeze this FrozenDict.

Returns

An unfrozen version of this FrozenDict instance.

Return type

Dict[flax.core.frozen_dict.K, flax.core.frozen_dict.V]

flax.core.frozen_dict.freeze(xs)

Freeze a nested dict.

Makes a nested dict immutable by transforming it into FrozenDict.

Parameters

xs (Mapping[Any, Any]) –

Return type

flax.core.frozen_dict.FrozenDict[Any, Any]

flax.core.frozen_dict.unfreeze(x)

Unfreeze a FrozenDict.

Makes a mutable copy of a FrozenDict mutable by transforming it into (nested) dict.

Parameters

x (flax.core.frozen_dict.FrozenDict[Any, Any]) –

Return type

Dict[Any, Any]