[OutputRecorder] re.search on layer_name

src/transformers/modeling_utils.py

@@ -1192,18 +1192,31 @@ def can_record_outputs(self) -> dict[str, OutputRecorder]:
         ```
 
          This means you can record outputs from the same class, by specifying a layer name. Before
-         collecting outputs, we check that they come from this layer.
+         collecting outputs, we check that they come from this layer. `layer_name` is a regex pattern
+         (matched with `re.search` against the submodule's dotted qualified name), so anchors can be used
+         to target a single index without prefix-matching siblings (e.g. `"layers\\.1$"` matches `layers.1`
+         but not `layers.10`).
 
          If you have cross attention that come from `LlamaAttention` and self attention that also
          come from `LlamaAttention` but from `self_attn` you can do this:
 
          ```python
          class LlamaModel(PreTrainedModel):
              _can_record_outputs = {
-                 "attentions": OutputRecorder(LlamaAttention, index=1, layer-name="self_attn"),
+                 "attentions": OutputRecorder(LlamaAttention, index=1, layer_name="self_attn"),
                  "cross_attentions": OutputRecorder(LlamaAttention, index=1, layer_name="cross_attn")
              }
 
+        ```
+
+         Regex alternation can also be used to pick a non-contiguous subset of layers, e.g. to
+         capture hidden states from layers 6, 12, and 18 only:
+
+         ```python
+         class MyModel(PreTrainedModel):
+             _can_record_outputs = {
+                 "hidden_states": OutputRecorder(MyBlock, layer_name=r"layers\\.(6|12|18)$"),
+             }
         ```
         """
         return self._can_record_outputs or {}

src/transformers/models/blip/modeling_blip_text.py

@@ -444,10 +444,10 @@ class BlipTextPreTrainedModel(PreTrainedModel):
     _can_record_outputs = {
         "hidden_states": BlipTextLayer,
         "attentions": [
-            OutputRecorder(BlipTextSelfAttention, index=1, layer_name=".attention."),
+            OutputRecorder(BlipTextSelfAttention, index=1, layer_name=r"\.attention\."),
         ],
         "cross_attentions": [
-            OutputRecorder(BlipTextSelfAttention, index=1, layer_name=".crossattention."),
+            OutputRecorder(BlipTextSelfAttention, index=1, layer_name=r"\.crossattention\."),
         ],
     }
 

src/transformers/models/blip_2/modeling_blip_2.py

@@ -887,10 +887,10 @@ class Blip2QFormerModel(Blip2PreTrainedModel):
     _can_record_outputs = {
         "hidden_states": Blip2QFormerLayer,
         "attentions": [
-            OutputRecorder(Blip2QFormerMultiHeadAttention, index=1, layer_name=".attention"),
+            OutputRecorder(Blip2QFormerMultiHeadAttention, index=1, layer_name=r"\.attention"),
         ],
         "cross_attentions": [
-            OutputRecorder(Blip2QFormerMultiHeadAttention, index=1, layer_name=".crossattention"),
+            OutputRecorder(Blip2QFormerMultiHeadAttention, index=1, layer_name=r"\.crossattention"),
         ],
     }
 

src/transformers/models/decision_transformer/modeling_decision_transformer.py

@@ -316,8 +316,8 @@ class DecisionTransformerGPT2PreTrainedModel(PreTrainedModel):
     _can_compile_fullgraph = False
     _can_record_outputs = {
         "hidden_states": DecisionTransformerGPT2Block,
-        "attentions": OutputRecorder(DecisionTransformerGPT2Attention, layer_name=".attn", index=1),
-        "cross_attentions": OutputRecorder(DecisionTransformerGPT2Attention, layer_name=".crossattention", index=1),
+        "attentions": OutputRecorder(DecisionTransformerGPT2Attention, layer_name=r"\.attn", index=1),
+        "cross_attentions": OutputRecorder(DecisionTransformerGPT2Attention, layer_name=r"\.crossattention", index=1),
     }
 
     # No longer used as we directly use our masks instead

src/transformers/models/gpt2/modeling_gpt2.py

@@ -422,8 +422,8 @@ class GPT2PreTrainedModel(PreTrainedModel):
     _can_compile_fullgraph = True
     _can_record_outputs = {
         "hidden_states": GPT2Block,
-        "attentions": OutputRecorder(GPT2Attention, layer_name=".attn", index=1),
-        "cross_attentions": OutputRecorder(GPT2Attention, layer_name=".crossattention", index=1),
+        "attentions": OutputRecorder(GPT2Attention, layer_name=r"\.attn", index=1),
+        "cross_attentions": OutputRecorder(GPT2Attention, layer_name=r"\.crossattention", index=1),
     }
 
     # No longer used as we directly use our masks instead

src/transformers/models/instructblip/modeling_instructblip.py

@@ -796,10 +796,10 @@ class InstructBlipQFormerModel(InstructBlipPreTrainedModel):
     _can_record_outputs = {
         "hidden_states": InstructBlipQFormerLayer,
         "attentions": [
-            OutputRecorder(InstructBlipQFormerMultiHeadAttention, index=1, layer_name=".attention"),
+            OutputRecorder(InstructBlipQFormerMultiHeadAttention, index=1, layer_name=r"\.attention"),
         ],
         "cross_attentions": [
-            OutputRecorder(InstructBlipQFormerMultiHeadAttention, index=1, layer_name=".crossattention"),
+            OutputRecorder(InstructBlipQFormerMultiHeadAttention, index=1, layer_name=r"\.crossattention"),
         ],
     }
 

src/transformers/models/instructblipvideo/modeling_instructblipvideo.py

@@ -745,10 +745,10 @@ class InstructBlipVideoQFormerModel(InstructBlipVideoPreTrainedModel):
     _can_record_outputs = {
         "hidden_states": InstructBlipVideoQFormerLayer,
         "attentions": [
-            OutputRecorder(InstructBlipVideoQFormerMultiHeadAttention, index=1, layer_name=".attention"),
+            OutputRecorder(InstructBlipVideoQFormerMultiHeadAttention, index=1, layer_name=r"\.attention"),
         ],
         "cross_attentions": [
-            OutputRecorder(InstructBlipVideoQFormerMultiHeadAttention, index=1, layer_name=".crossattention"),
+            OutputRecorder(InstructBlipVideoQFormerMultiHeadAttention, index=1, layer_name=r"\.crossattention"),
         ],
     }
 

src/transformers/models/minimax/modeling_minimax.py

@@ -596,7 +596,7 @@ class MiniMaxPreTrainedModel(PreTrainedModel):
     _can_compile_fullgraph = False  # uses a non-compilable custom cache class MiniMaxCache
     _supports_attention_backend = True
     _can_record_outputs = {
-        "router_logits": OutputRecorder(MiniMaxTopKRouter, layer_name="mlp.gate", index=0),
+        "router_logits": OutputRecorder(MiniMaxTopKRouter, layer_name=r"mlp\.gate", index=0),
         "hidden_states": MiniMaxDecoderLayer,
         "attentions": [MiniMaxAttention, MiniMaxLightningAttention],
     }

src/transformers/models/minimax/modular_minimax.py

@@ -403,7 +403,7 @@ def forward(
 class MiniMaxPreTrainedModel(MixtralPreTrainedModel):
     _can_compile_fullgraph = False  # uses a non-compilable custom cache class MiniMaxCache
     _can_record_outputs = {
-        "router_logits": OutputRecorder(MiniMaxTopKRouter, layer_name="mlp.gate", index=0),
+        "router_logits": OutputRecorder(MiniMaxTopKRouter, layer_name=r"mlp\.gate", index=0),
         "hidden_states": MiniMaxDecoderLayer,
         "attentions": [MiniMaxAttention, MiniMaxLightningAttention],
     }

src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py

@@ -1047,7 +1047,7 @@ class Qwen3OmniMoeVisionEncoder(Qwen3OmniMoePreTrainedModel):
     input_modalities = ("image", "video")
     _no_split_modules = ["Qwen3OmniMoeVisionBlock"]
     _can_record_outputs = {
-        "router_logits": OutputRecorder(Qwen3OmniMoeTextTopKRouter, layer_name="mlp.gate", index=0),
+        "router_logits": OutputRecorder(Qwen3OmniMoeTextTopKRouter, layer_name=r"mlp\.gate", index=0),
         "hidden_states": Qwen3OmniMoeVisionBlock,
         "attentions": Qwen3OmniMoeVisionAttention,
     }

src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py

@@ -601,7 +601,7 @@ class Qwen3VLMoeVisionModel(Qwen3VLMoePreTrainedModel):
     input_modalities = ("image", "video")
     _no_split_modules = ["Qwen3VLMoeVisionBlock"]
     _can_record_outputs = {
-        "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name="mlp.gate", index=0),
+        "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name=r"mlp\.gate", index=0),
         "hidden_states": Qwen3VLMoeVisionBlock,
         "attentions": Qwen3VLMoeVisionAttention,
     }

src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py

@@ -228,7 +228,7 @@ class Qwen3VLMoeVisionBlock(Qwen3VLVisionBlock):
 
 class Qwen3VLMoeVisionModel(Qwen3VLVisionModel):
     _can_record_outputs = {
-        "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name="mlp.gate", index=0),
+        "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name=r"mlp\.gate", index=0),
         "hidden_states": Qwen3VLMoeVisionBlock,
         "attentions": Qwen3VLMoeVisionAttention,
     }

src/transformers/models/switch_transformers/modeling_switch_transformers.py

@@ -618,8 +618,8 @@ def _shift_right(self, input_ids):
 class SwitchTransformersStack(SwitchTransformersPreTrainedModel):
     _can_record_outputs = {
         "hidden_states": SwitchTransformersBlock,
-        "attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name="layer.0"),
-        "cross_attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name="layer.1"),
+        "attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name=r"layer\.0"),
+        "cross_attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name=r"layer\.1"),
         "router_logits": OutputRecorder(SwitchTransformersTop1Router, index=2),
     }
 

src/transformers/models/switch_transformers/modular_switch_transformers.py

@@ -398,8 +398,8 @@ def _shift_right(self, input_ids):
 class SwitchTransformersStack(SwitchTransformersPreTrainedModel):
     _can_record_outputs = {
         "hidden_states": SwitchTransformersBlock,
-        "attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name="layer.0"),
-        "cross_attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name="layer.1"),
+        "attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name=r"layer\.0"),
+        "cross_attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name=r"layer\.1"),
         "router_logits": OutputRecorder(SwitchTransformersTop1Router, index=2),
     }
 

src/transformers/models/vjepa2/modeling_vjepa2.py

@@ -861,8 +861,8 @@ class VJEPA2PreTrainedModel(PreTrainedModel):
     _supports_sdpa = True
     _supports_flash_attn = True
     _can_record_outputs = {
-        "hidden_states": OutputRecorder(VJEPA2Layer, layer_name="encoder.layer"),
-        "attentions": OutputRecorder(VJEPA2RopeAttention, index=1, layer_name="encoder.layer"),
+        "hidden_states": OutputRecorder(VJEPA2Layer, layer_name=r"encoder\.layer"),
+        "attentions": OutputRecorder(VJEPA2RopeAttention, index=1, layer_name=r"encoder\.layer"),
     }
 
     @torch.no_grad()

src/transformers/utils/output_capturing.py

@@ -18,6 +18,7 @@
 
 from __future__ import annotations
 
+import re
 import threading
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -39,13 +40,19 @@
 @dataclass
 @requires(backends=("torch",))
 class OutputRecorder:
-    """
+    r"""
     Configuration for recording outputs from a model via hooks.
 
     Attributes:
         target_class (Type): The class (e.g., nn.Module) to which the hook will be attached.
         index (Optional[int]): If the output is a tuple/list, optionally record only at a specific index.
-        layer_name (Optional[str]): Name of the submodule to target (if needed), e.g., "transformer.layer.3.attn".
+        layer_name (Optional[str]): Regex pattern (matched with `re.search`) used to filter submodules by their
+            dotted qualified name. Examples:
+
+            - `"self_attn"`: substring match
+            - `"transformer.layer.3.attn"`: literal path
+            - `r"layers\.1$"`: anchored, targets layer 1 without also matching `layers.10`/`layers.11`/…
+            - `r"layers\.(6|12|18)$"`: picks a non-contiguous subset of layers
         class_name (Optional[str]): Name of the class to which the hook will be attached. Could be the suffix of class name in some cases.
     """
 
@@ -142,7 +149,7 @@ def recursively_install_hooks(
         if (specs.target_class is not None and isinstance(parent_module, specs.target_class)) or (
             specs.class_name is not None and module_name.endswith(specs.class_name)
         ):
-            if specs.layer_name is not None and specs.layer_name not in module_name:
+            if specs.layer_name is not None and re.search(specs.layer_name, module_name) is None:
                 continue
             install_output_capturing_hook(parent_module, key, specs.index)