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.

pop(key)

Create a new FrozenDict where one entry is removed.

Example:

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

Parameters

key – the key to remove from the dict

Returns

A pair with the new FrozenDict and the removed value.

pretty_repr(num_spaces=4)

Returns an indented representation of the nested dictionary.

unfreeze()

Unfreeze this FrozenDict.

Returns

An unfrozen version of this FrozenDict instance.

flax.core.frozen_dict.freeze(xs)

Freeze a nested dict.

Makes a nested dict immutable by transforming it into FrozenDict.

Parameters

xs – Dictionary to freeze (a regualr Python dict).

Returns

The frozen dictionary.

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 – Frozen dictionary to unfreeze.

Returns

The unfrozen dictionary (a regular Python dict).

flax.core.frozen_dict.copy(x, add_or_replace)

Create a new dict with additional and/or replaced entries. This is a utility function that can act on either a FrozenDict or regular dict and mimics the behavior of FrozenDict.copy.

Example:

new_variables = copy(variables, {‘additional_entries’: 1})

Parameters

• x – the dictionary to be copied and updated

• add_or_replace – dictionary of key-value pairs to add or replace in the dict x

Returns

A new dict with the additional and/or replaced entries.

flax.core.frozen_dict.pop(x, key)

Create a new dict where one entry is removed. This is a utility function that can act on either a FrozenDict or regular dict and mimics the behavior of FrozenDict.pop.

Example:

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

Parameters

• x – the dictionary to remove the entry from

• key – the key to remove from the dict

Returns

A pair with the new dict and the removed value.