t5x.config_utils package

Utilities for configuring T5X binaries.

t5x.config_utils.config_with_fiddle(function)

Configure and build a T5X launcher from Fiddle command line flags.

The output config, when called via fdl.build(), will execute function.

Parameters:

function – A function that launches a T5X job, e.g., train, eval, …

Returns:

The buildable of the function or object, depending on whether

–fdl_config_file or –fdl_config was passed.

Raises:

• AssertionError – If not using_fdl().

• ValueError – If both fiddle and gin arguments were passed on the command line.

• ValueError – If both –fdl_config_file and –fdl_config were passed.

• ValueError – If the object built via –fdl_config does not build as a call to function.

t5x.config_utils.direct_summarize_fiddle_config(model_dir, summary_writer, step, get_current_fiddle_config=None)

Writes fiddle config to the model dir and TensorBoard summary.

When passing this function to your fiddle config, don’t pass the private function; instead pass the fdl.Partial version of summarize_fiddle_config.

Parameters:

• model_dir – Model directory to write to.

• summary_writer – MetricWriter, if any.

• step – Current step.

• get_current_fiddle_config – This will be filled in by t5x.config_utils.prepare_to_summarize().

t5x.config_utils.prepare_to_summarize(config)

Update config so its summarize_fiddle_config calls can access it.

t5x.config_utils.run(main)

Wrapper for app.run that rewrites jax, gin, and fiddle flags.

t5x.config_utils.sanitize_summary_getter(config)

Update config to remove get_current_fiddle_config calls.

t5x.config_utils.using_fdl()

Returns true if any Fiddle configuration flags are set.