Gets the shared module and state with the given name from a nearest ancestor. Shared modules should be registered via `_share_with_descendants`. Args: shared_module_name: The name of the shared module. Returns: The shared module and corresponding st
(self, shared_module_or_name: Union["Module", str])
| 1024 | f"Shared module already exists under name {shared_module_name}: " |
| 1025 | f"{self._paths_to_shared_modules[shared_module_name]} vs. {relative_path}" |
| 1026 | ) |
| 1027 | self._paths_to_shared_modules[shared_module_name] = relative_path |
| 1028 | self._shared_module_names[module] = shared_module_name |
| 1029 | |
| 1030 | class SharedModuleInfo(NamedTuple): |
| 1031 | name: str |
| 1032 | module: "Module" |
| 1033 | state: NestedTensor |
| 1034 | |
| 1035 | def get_shared_module(self, shared_module_or_name: Union["Module", str]) -> SharedModuleInfo: |
| 1036 | """Gets the shared module and state with the given name from a nearest ancestor. |
| 1037 | |
| 1038 | Shared modules should be registered via `_share_with_descendants`. |
| 1039 | |
| 1040 | Args: |
| 1041 | shared_module_name: The name of the shared module. |
| 1042 | |
| 1043 | Returns: |
| 1044 | The shared module and corresponding state. |
| 1045 | |
| 1046 | Raises: |
| 1047 | ValueError: if `shared_module_name` has not been shared by any ancestor. |
| 1048 | """ |
| 1049 | # pylint: disable=protected-access |
| 1050 | context = self.get_invocation_context() |
| 1051 | |
| 1052 | def context_shares_module( |
| 1053 | ctx: InvocationContext, |
| 1054 | ) -> bool: # pytype: disable=invalid-annotation |
| 1055 | if isinstance(shared_module_or_name, str): |
| 1056 | return shared_module_or_name in ctx.module._paths_to_shared_modules |
| 1057 | elif isinstance(shared_module_or_name, Module): |
| 1058 | return shared_module_or_name in ctx.module._shared_module_names |
| 1059 | raise ValueError(f"{shared_module_or_name=} must be a string or Module.") |
| 1060 | |
| 1061 | while context is not None and not context_shares_module(context): |
| 1062 | context = context.parent |
| 1063 | if context is None: |
| 1064 | raise InvalidDescendantError( |
| 1065 | f"Module '{self.path()}' does not have an ancestor that shares " |
| 1066 | f"{shared_module_or_name=}." |
| 1067 | ) |
| 1068 | |
| 1069 | if isinstance(shared_module_or_name, Module): |
| 1070 | shared_module_or_name = context.module._shared_module_names[shared_module_or_name] |
| 1071 | assert isinstance(shared_module_or_name, str) |
| 1072 | |
| 1073 | target_module, target_state = context.module, context.state |
| 1074 | # pylint: disable-next=protected-access |
| 1075 | path_from_ancestor = context.module._paths_to_shared_modules[shared_module_or_name] |
| 1076 | for part in path_from_ancestor: |
| 1077 | if part not in target_module.children: |
| 1078 | raise InvalidDescendantError( |
| 1079 | f"Module '{target_module.path()}' does not contain '{part}' from path " |
| 1080 | f"'{'.'.join(path_from_ancestor)}'" |
| 1081 | ) |
| 1082 | if part not in target_state: |
| 1083 | raise InvalidDescendantError( |