flax.config package

Global configuration options for Flax.

Now a wrapper over jax.config, in which all config vars have a ‘flax_’ prefix.

To modify a config value on run time, call: flax.config.update('flax_<config_name>', <value>)

flax.configurations.define_bool_state(name, default, help)

Set up a boolean flag using JAX’s config system.

The flag will actually be stored as an environment variable of ‘FLAX_<UPPERCASE_NAME>’. JAX config ensures that the flag can be overwritten on runtime with flax.config.update(‘flax_<config_name>’, <value>).

flax.configurations.static_bool_env(varname, default)

Read an environment variable and interpret it as a boolean.

This is deprecated. Please use define_bool_state() unless your flag will be used in a static method and does not require runtime updates.

True values are (case insensitive): ‘y’, ‘yes’, ‘t’, ‘true’, ‘on’, and ‘1’; false values are ‘n’, ‘no’, ‘f’, ‘false’, ‘off’, and ‘0’. :param varname: the name of the variable :param default: the default boolean value

Returns

boolean return value derived from defaults and environment.

Raises: ValueError if the environment variable is anything else.