Method get_invocation_context

axlearn/common/module.py:1084–1096

(self)

            if part not in target_state:
                raise InvalidDescendantError(
                    f"Module '{target_module.path()}' state does not contain '{part}' from path "
                    f"'{'.'.join(path_from_ancestor)}'. The state contains: {target_state.keys()}"
                )
            target_module, target_state = target_module.children[part], target_state[part]

        return Module.SharedModuleInfo(
            module=target_module, state=target_state, name=shared_module_or_name
        )

    def get_invocation_context(self) -> InvocationContext:  # pytype: disable=invalid-annotation
        context = current_context()
        if not context:
            raise RuntimeError(
                "Module invocation context not found. "
                "Did you invoke the module inside functional(...)?"
            )