Function per_param_setting_debug_string

axlearn/experiments/test_utils.py:105–143

Given module and config, returns a dictionary of read_per_param_settings debug strings. If no register_per_param_settings is called, an assertion error is raised.

(
    module: str,
    config_name: str,
    trainer_config: TrainerConfigFn,
    *,
    kv_separator=": ",
    field_separator="\n",
)

def per_param_setting_debug_string(
    module: str,
    config_name: str,
    trainer_config: TrainerConfigFn,
    *,
    kv_separator=": ",
    field_separator="\n",
) -> dict[str, str]:
    """Given module and config, returns a dictionary of read_per_param_settings debug strings.

    If no register_per_param_settings is called, an assertion error is raised.
    """
    all_param_settings = read_per_param_settings(
        module=module, config_name=config_name, trainer_config=trainer_config
    )
    if not all_param_settings:
        raise AssertionError(
            f"{module}.{config_name} does not use per_param_setting_by_path."
            "It is safe to remove this test."
        )

    def sep_line(name: str):
        return "=" * 20 + name + "=" * 20

    def _settings_to_string(description: str, settings: dict[str, NestedTree]) -> str:
        lines = []

        for setting_key, setting_value in settings.items():
            # If there are multiple settings under the same description key,
            # separate settings generated by different optimizers/learners by sep_line.
            lines.append(sep_line(name=f"{description} {setting_key}"))
            for param_name, param_setting in flatten_items(setting_value):
                lines.append(f"{param_name}{kv_separator}{param_setting}")
        return field_separator.join(lines)

    return {
        description: _settings_to_string(description, settings)
        for description, settings in all_param_settings.items()
    }


@enum.unique