Class Module

axlearn/common/module.py:772–1148

A node in a tree of Modules.

            kwargs_types = {k: type(v) for k, v in kwargs.items()}
            new_exc = TypeError(
                f"Type error when calling {self}.{method_fn} "
                f"with args={args_types} and kwargs={kwargs_types}"
            )
            setattr(new_exc, "_handled", True)
            raise new_exc from e

    return wrap_method_fn


class Module(Configurable, metaclass=_PostInitMeta):
    """A node in a tree of Modules."""

    @config_class
    class Config(Configurable.Config):
        """Module config.

        Attributes:
            name: Name of this module.
            vlog: The maximum vlog level. If None, vlog is disabled.
        """

        name: Required[str] = REQUIRED
        vlog: Optional[int] = None

    def __init__(self, cfg: Config, *, parent: Optional["Module"]):
        super().__init__(cfg)
        cfg = self.config
        self._name = cfg.name
        self._parent = parent  # Avoid adding parent to self._modules.
        self._children: dict[str, "Module"] = {}
        # Mapping from descendant module name to relative path from current module.
        self._paths_to_shared_modules: dict[str, list[str]] = {}
        # Mapping from modules being shared by the current module, to the shared module name.
        self._shared_module_names: dict["Module", str] = {}
        self._vlog_level = cfg.vlog

    def __post_init__(self):
        # Wrap methods after `__init__`, allowing access to child modules.
        for method_name, method_fn in self._wrapped_methods_for_auto_child_context().items():
            setattr(self, method_name, method_fn)

    def _wrapped_methods_for_auto_child_context(self) -> dict[str, Callable]:
        """Returns methods that have been wrapped and bound to `self`.

        This ensures that module methods are bound to the instance that defined the method, rather
        than the instance that the method is assigned to in `__post_init__`.

        For example, `self.child._wrapped_methods_for_auto_child_context()` returns methods bound to
        `self.child` rather than `self`, which affects what `self.config` points to within the
        wrapped method.

        On the other hand, `self.child._methods_to_wrap_for_auto_child_context()` returns un-bound
        methods of `self.child`. Subclasses will typically override this method to control which
        methods of the subclass to wrap.
        """
        methods = self._methods_to_wrap_for_auto_child_context()
        return self._wrap_methods_with_auto_child_context(methods)