flax.config package#

Flax has the following global config options.

Global configuration options for Flax.

flax.config.flax_filter_frames#

Whether to hide flax-internal stack frames from tracebacks. Set by the FLAX_FILTER_FRAMES environment variable. Defaults to True.

flax.config.flax_profile#

Whether to automatically wrap Module methods with jax.named_scope for profiles. Set by the FLAX_PROFILE environment variable. Defaults to True.