AutoConfig: Reinterpret builder functions for deep configurability

The @auto_config function decorator reinterprets the function to construct a data structure of fdl.Config’s that correspond to the objects instead of the objects themselves.

auto_config

fiddle.experimental.auto_config.auto_config(fn=None, *, experimental_allow_dataclass_attribute_access=False, experimental_allow_control_flow=False, experimental_always_inline=None, experimental_exemption_policy=None, experimental_config_types=ConfigTypes(config_cls=<class 'fiddle._src.config.Config'>, partial_cls=<class 'fiddle._src.partial.Partial'>, arg_factory_cls=<class 'fiddle._src.partial.ArgFactory'>), experimental_result_must_contain_buildable=True)

Rewrites the given function to make it generate a Config.

This function creates a new function from fn by rewriting its AST (abstract syntax tree), replacing all Call nodes with a custom call handler. When the rewritten function is run, the call handler intercepts calls and applies the following rules:

• Calls to builtins, methods, callables without an inferrable signature, callables wrapped by auto_config.exempt, or other functions that have been auto_config’ed take place as usual.

• Calls to functools.partial are replaced by calling fdl.Partial with the same arguments;

• All other calls are replaced by calling fdl.Config with the arguments that would have been passed to the called function or class.

This function may be used standalone or as a decorator. The returned function is simply a wrapper around fn, but with an additional as_buildable attribute containing the rewritten function. For example:

def build_model():
  return Sequential([
      Dense(num_units=128, activation=relu),
      Dense(num_units=128, activation=relu),
      Dense(num_units=1, activation=None),
  ])

config = auto_config(build_model).as_buildable()

The resulting config is equivalent to the following “manually” constructed configuration graph:

fdl.Config(Sequential, layers=[
    fdl.Config(Dense, num_units=128, activation=relu),
    fdl.Config(Dense, num_units=128, activation=relu),
    fdl.Config(Dense, num_units=1, activation=None),
])

This can then be built with fdl.build(config). Without modification, this will result in the same model as just calling build_model() directly. However, config permits changes to the model hyperparameters, for example:

config.layers[0].num_units = 64
config.layers[0].activation = 'elu'
config.layers[1].num_units = 64
config.layers[1].activation = 'elu'

modified_model = fdl.build(config)

Currently, control flow is not supported by default in auto_config. Experimental support for control flow can be enabled using the experimental_allow_control_flow argument. If enabled, control flow constructs may be used within the function to construct the resulting config (for example, a for loop could be used to build a list of layers). Control flow is never encoded directly as part of the resulting fdl.Config (for example, there is no fdl.Config that will correspond to a conditional or loop). While many simple constructs (for _ in range(10) etc) work, there will also likely be surprising behavior in some circumstances (for example, using itertools functions in conjunction with a loop will not work, since the calls to itertools functions will be turned into fdl.Config objects).

Using @auto_config is compatible with both @staticmethod and @classmethod, however the @auto_config decorator must appear above the @classmethod or @staticmethod in the decorator list.

Parameters:

• fn – The function to create a config-generating function from.

• experimental_allow_dataclass_attribute_access – Whether to allow attribute access on dataclasses within auto_config. Note that access to dataclass attribute is transformed into access to fdl.Config attributes in the as_buildable path.

• experimental_allow_control_flow (bool) – Whether to allow control flow constructs in fn. By default, control flow constructs will cause an UnsupportedLanguageConstructError to be thrown.

• experimental_always_inline (Optional[bool]) – If true, this function (when called in an auto_config context) will always be inline’d in-place. See the documentation on inline for an example. The default (if unspecified) is currently True.

• experimental_exemption_policy (Optional[Callable[[Type[Any]], bool]]) – An optional policy to control which function calls within the body of fn should be turned into fdl.Config’s and which ones should simply be executed normally during the as_buildable interpretation of fn. This predicate should return True if the given callable should be exempted from auto-configuration.

• experimental_config_types (ConfigTypes) – A ConfigTypes instance containing the types to use when generating configs. By default, this just supplies the standard Fiddle types ()``fdl.Config``, fdl.Partial, and fdl.ArgFactory), but projects with custom subclasses can use this to override the default. This is experimental and may be removed in the future.

• experimental_result_must_contain_buildable (bool) – If true, then raise an error if fn.as_buildable returns a result that does not contain any Buildable values – e.g., if it returns an empty dict.

Return type:

Any

Returns:

A wrapped version of fn, but with an additional as_buildable attribute containing the rewritten function.

is_auto_config

fiddle.experimental.auto_config.is_auto_config(function_object)

Return type:

bool

auto_unconfig

fiddle.experimental.auto_config.auto_unconfig(fn=None, *, experimental_always_inline=None)

Converts functions that create buildables directly into auto_config form.

While most of the time, the benefits of an auto_config representation of object configuration and construction are valuable (e.g. static type checking and tooling / refactoring support), sometimes it is more convenient to manipulate buildable objects directly. auto_unconfig converts a function that directly manipulates fdl.Buildable’s (e.g. fdl.Config’s) into one that looks identically to an auto_config’d function, and is fully interoperable with the rest of the auto_config ecosystem.

Example:

@auto_unconfig
def make_experiment_trainer(name: str) -> fdl.Config[MyTrainer]
  model = make_model.as_buildable(name)
  select(model, DropOut).set(rate=0.42)  # Full Fiddle API available!
  dataset = make_dataset.as_buildable()
  # Build fdl.Config's imperatively.
  trainer_config = fdl.Config(MyTrainer)
  trainer_config.model = model
  trainer_config.train_dataset = dataset
  trainer_config.skip_eval = True
  return trainer_config  # Return a `fdl.Buildable`

# Sample usage within an auto_config'd function.
@auto_config
def make_driver():
  return TrainerDriver(
    trainer=make_experiment_trainer('my_experiment'),
    checkpointer=CustomCheckpointer())

# Sample usage outside of auto_config contexts.
def main():
  # Use instantiated objects:
  trainer = make_experiment_trainer('my_experiment')
  for example in trainer.train_dataset:
    print_prediction(trainer.model.predict(example))

  # Or manipulate the configuration before calling `fdl.build`:
  trainer_config = make_experiment_trainer.as_buildable('my_experiment')
  trainer_config.skip_eval = False  # Tweak configuration.
  trainer2 = fdl.build(trainer_config)
  run_trainer(trainer2)

Parameters:

• fn – The function to convert.

• experimental_always_inline (Optional[bool]) – Whether the output of fn should always be inlined into the caller’s config. The default (if unspecified) is True.

Return type:

Any

Returns:

An AutoConfig that corresponds to fn.

inline

fiddle.experimental.auto_config.inline(buildable)

Converts an auto_config-based buildable into a DAG of Buildables.

inline updates buildable in place to preserve aliasing within a larger Fiddle configuration. If you would like to leave buildable unmodified, make a shallow copy (copy.copy) before calling inline.

Example:

# shared/input_pipelines.py
@auto_config(experimental_always_inline=False)
def make_input_pipeline(name: str, batch_size: int) -> InputPipeline:
  file_path = '/base_path/'+name
  augmentation = 'my_augmentation_routine'
  # ...
  return InputPipeline(file_path, augmentation, ...)

# config/main.py
@auto_config
def make_experiment():
  data = make_input_pipeline('normal_dataset', batch_size)
  model = ...
  return Experiment(data, model)

# experiment_configuration.py
def make_experiment():
  config = make_experiment.as_buildable()
  config.data.name = 'advanced_dataset'
  # config.data.augmentation = 'custom_augmentation'  # Not configurable!!!
  # return fdl.build(config)                          # Works like normal.
  auto_config.inline(config.data)
  print(config.data.file_path)       # Prints: '/base_path/advanced_dataset'
  config.data.augmentation = 'custom_augmentation'    # Now exposed.
  experiment = fdl.build(config)                      # Works like normal.
  return experiment

Parameters:

buildable (Config) – The buildable of an auto_config’d function to replace with the root of a Fiddle DAG that corresponds to it.

Raises:

ValueError – If buildable is not a Config, or if buildable doesn’t correspond to an auto_config’d function.

AutoConfig

class fiddle.experimental.auto_config.AutoConfig(func, buildable_func, always_inline)

A function wrapper for auto_config’d functions.

In order to support auto_config’ing @classmethod’s, we need to customize the descriptor protocol for the auto_config’d function. This simple wrapper type is designed to look like a functool.wraps wrapper, but implements custom behavior for bound methods.