flax.config package

Global configuration flags for Flax.

class flax.configurations.Config

property flax_filter_frames

Whether to hide flax-internal stack frames from tracebacks.

property flax_fix_rng_separator

Whether to add separator characters when folding in static data into PRNG keys.

property flax_preserve_adopted_names

When adopting outside modules, don’t clobber existing names.

property flax_profile

Whether to run Module methods under jax.named_scope for profiles.

property flax_return_frozendict

Whether to return FrozenDicts when calling init or apply.

property flax_use_orbax_checkpointing

Whether to use Orbax to save checkpoints.

update(name_or_holder, value, /)

Modify the value of a given flag.

Parameters

• name_or_holder – the name of the flag to modify or the corresponding flag holder object.

• value – new value to set.