Function wrapper

axlearn/experiments/trainer_config_utils.py:205–216

(*args, **kwargs)

    @wraps(func)
    def wrapper(*args, **kwargs) -> dict[str, TrainerConfigFn]:
        result = cached_func(get_data_dir(), *args, **kwargs)
        new_result = {}
        for key, fn in result.items():
            if getattr(fn, "_is_wrapped_with_deepcopy", False):
                # No need to wrap the function again. This can happen when a parent module
                # recursively combines other named_trainer_configs output from its submodules.
                new_result[key] = fn
            else:
                # Wrap it so that a true deepcopy with closures will happen first.
                new_result[key] = _wrap_with_deep_copy_with_closure(fn)
        return new_result

    return wrapper