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)[source]#
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)[source]#
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.