flax.core.frozen_dict package#

class flax.core.frozen_dict.FrozenDict(*args, __unsafe_skip_copy__=False, **kwargs)[source]#

An immutable variant of the Python dict.

copy(add_or_replace=mappingproxy({}))[source]#

Create a new FrozenDict with additional or replaced entries.

Parameters:

add_or_replace (Mapping[TypeVar(K), TypeVar(V)]) –

Return type:

FrozenDict[TypeVar(K), TypeVar(V)]

pop(key)[source]#

Create a new FrozenDict where one entry is removed.

Example:

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

key (TypeVar(K)) – the key to remove from the dict

Return type:

Tuple[FrozenDict[TypeVar(K), TypeVar(V)], TypeVar(V)]

Returns:

A pair with the new FrozenDict and the removed value.

pretty_repr(num_spaces=4)[source]#

Returns an indented representation of the nested dictionary.

unfreeze()[source]#

Unfreeze this FrozenDict.

Return type:

Dict[TypeVar(K), TypeVar(V)]

Returns:

An unfrozen version of this FrozenDict instance.

flax.core.frozen_dict.freeze(xs)[source]#

Freeze a nested dict.

Makes a nested dict immutable by transforming it into FrozenDict.

Parameters:

xs (Mapping[Any, Any]) – Dictionary to freeze (a regualr Python dict).

Return type:

FrozenDict[Any, Any]

Returns:

The frozen dictionary.

flax.core.frozen_dict.unfreeze(x)[source]#

Unfreeze a FrozenDict.

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

Parameters:

x (Union[FrozenDict, Dict[str, Any]]) – Frozen dictionary to unfreeze.

Return type:

Dict[Any, Any]

Returns:

The unfrozen dictionary (a regular Python dict).

flax.core.frozen_dict.copy(x, add_or_replace=FrozenDict({}))[source]#

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:
Return type:

Union[FrozenDict, Dict[str, Any]]

Returns:

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

flax.core.frozen_dict.pop(x, key)[source]#

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:
Return type:

Tuple[Union[FrozenDict, Dict[str, Any]], Any]

Returns:

A pair with the new dict and the removed value.

flax.core.frozen_dict.pretty_repr(x, num_spaces=4)[source]#

Returns an indented representation of the nested dictionary. This is a utility function that can act on either a FrozenDict or regular dict and mimics the behavior of FrozenDict.pretty_repr. If x is any other dtype, this function will return repr(x).

Parameters:
  • x (Any) – the dictionary to be represented

  • num_spaces (int) – the number of space characters in each indentation level

Return type:

str

Returns:

An indented string representation of the nested dictionary.