flax.config package#

Global configuration flags for Flax.

class flax.configurations.Config[source]#
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, /)[source]#

Modify the value of a given flag.

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

  • value – new value to set.