Repository: huggingface/transformers
File: src/transformers/core_model_loading.py (Python)

Code before:

    def __post_init__(self):
        if isinstance(self.source_patterns, str):
            self.source_patterns = [self.source_patterns]
        if isinstance(self.target_patterns, str):
            self.target_patterns = [self.target_patterns]

        # Due to how our `_checkpoint_conversion_mapping` mappings are written, we need a few exceptions here
        # when instantiating the reverse mapping (i.e. the targets become sources, and sources become targets)
        # The issues lie in the sources usually, so here we need to check the targets for the reversed mapping

        # Process target_patterns: detect capturing groups and replace with \1
        # Store the original capturing group patterns for reverse mapping
        target_capturing_groups: list[str | None] = []
        for i, pattern in enumerate(self.target_patterns):
            self.target_patterns[i], captured_group = reverse_target_pattern(pattern)
            if captured_group is not None:
                target_capturing_groups.append(captured_group)

        # We also need to check capturing groups in the sources during reverse mapping (e.g. timm_wrapper, sam3)
        capturing_groups_index = 0
        for i, pattern in enumerate(self.source_patterns):
            if r"\1" in pattern:
                # Use the stored capturing group from target_patterns
                pattern = pattern.replace(r"\1", target_capturing_groups[capturing_groups_index], 1)
                capturing_groups_index += 1
            self.source_patterns[i] = pattern

        # Construct the regex we will use to rename keys from the sources to the targets
        branches = []
        for i, source_pattern in enumerate(self.source_patterns):
            group_name = f"g{i}"
            pattern = source_pattern.replace(".*.", r"\..*\.")
            branches.append(f"(?P<{group_name}>{pattern})")
        self.compiled_sources = re.compile("|".join(branches))

    def add_tensor(self, target_key: str, source_key: str, source_pattern: str, future: Future):
        self.collected_tensors[source_pattern].append(future)
        self.layer_targets[target_key].add(source_key)

    def rename_source_key(self, source_key: str) -> tuple[str, str | None]:
        """
        Return a tuple (renamed_key, source_pattern_producing_the_match).
        Try renaming `source_key` according to the source and target patterns of the current WeightTransform.
        In case of a one-to-many transform, i.e. we have several target patterns, the matching source pattern
        will be replaced by the first of all the target patterns (they are then correctly expanded in the Operations).
        """
        # Try matching one of the alternation branches
        match_object = self.compiled_sources.search(source_key)
        if match_object is None:
            return source_key, None

Review comment (suggestion):
Here in all generality a purely index-based matching will break with a Transform with different number of `sources` and `targets` (e.g. `Operation(sources=r"blabla.(\d+)", targets=["bla1.\1", "bla2.\1"])` will break during reverse op, even though it is correctly defined in terms of patterns replacement as we have only 1 captured pattern)

If we have several sources with different capturing groups, then this is much harder to correctly match back in general (I believe it's actually non-determined in general as we cannot know what to match to what)

But so TLDR, let's make sure we only have 1 pattern inside `target_capturing_groups` at the end, and raise if we have several for now, and always replace with this unique pattern on ALL the sources

Code after:
    layer_targets: dict[str, set[str]] = field(default_factory=lambda: defaultdict(set), init=False)

    def __post_init__(self):
        if isinstance(self.source_patterns, str):
            self.source_patterns = [self.source_patterns]
        if isinstance(self.target_patterns, str):
            self.target_patterns = [self.target_patterns]

        # Due to how our `_checkpoint_conversion_mapping` mappings are written, we need a few exceptions here
        # when instantiating the reverse mapping (i.e. the targets become sources, and sources become targets)
        # The issues lie in the sources usually, so here we need to check the targets for the reversed mapping

        # Process target_patterns: detect capturing groups and replace with \1
        # Store the original capturing group patterns for reverse mapping
        target_capturing_groups: list[str] = []
        for i, pattern in enumerate(self.target_patterns):
            self.target_patterns[i], captured_group = process_target_pattern(pattern)
            if captured_group is not None:
                target_capturing_groups.append(captured_group)

        # Validate that we only have one unique capturing group pattern across all targets
        # This ensures deterministic reverse mapping when sources have \1 backreferences
        unique_capturing_groups = set(target_capturing_groups)
        if len(unique_capturing_groups) > 1:
            raise ValueError(
                f"Multiple different capturing groups found in target_patterns: {unique_capturing_groups}. "
                f"All target patterns must use the same capturing group pattern."
            )
        unique_capturing_group = unique_capturing_groups.pop() if unique_capturing_groups else None

        # We also need to check capturing groups in the sources during reverse mapping (e.g. timm_wrapper, sam3)
        for i, pattern in enumerate(self.source_patterns):
            if r"\1" in pattern:
                if unique_capturing_group is None:
                    raise ValueError(
                        f"Source pattern '{pattern}' contains \\1 backreference, but no capturing groups "
                        f"found in target_patterns."
                    )
                # Use the unique capturing group from target_patterns for all sources
                pattern = pattern.replace(r"\1", unique_capturing_group, 1)
            self.source_patterns[i] = pattern

        # Construct the regex we will use to rename keys from the sources to the targets
        branches = []
        for i, source_pattern in enumerate(self.source_patterns):
            group_name = f"g{i}"
            pattern = source_pattern.replace(".*.", r"\..*\.")
            branches.append(f"(?P<{group_name}>{pattern})")
        self.compiled_sources = re.compile("|".join(branches))

    def add_tensor(self, target_key: str, source_key: str, source_pattern: str, future: Future):