transformers/utils/rules.toml

# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This file can carry repo-local rule overrides for faster iteration between
# `transformers-mlinter` releases.
# Keep it synced with the upstream package's rules.toml when possible so local
# behavior does not drift from the published checker longer than necessary.

version = 1

[rules.TRF001]
description = "config_class on <Model>PreTrainedModel should match <Model>Config naming."
default_enabled = true
allowlist_models = ["qwen3_omni_moe"]

[rules.TRF001.explanation]
what_it_does = "Checks that <Model>PreTrainedModel's config_class is named <Model>Config."
why_bad = "A mismatch can break loading, auto classes and developer expectations."
diff = '''
 class AcmePreTrainedModel(PreTrainedModel):
-    config_class = WileConfig
+    config_class = AcmeConfig
'''

[rules.TRF002]
description = "base_model_prefix on a PreTrainedModel, when set, should be a non-empty canonical string."
default_enabled = true
# `lighton_ocr` sets `base_model_prefix = ""` on `LightOnOcrModel` (`modular_lighton_ocr.py`). It does not
# currently fire: the rule only inspects classes it can resolve to `PreTrainedModel`, and it resolves base
# chains within a single file, so `LightOnOcrModel(Mistral3Model)` is not recognized as a pretrained class in
# the modular file. Kept because the empty prefix is real and authored here, not inherited -- the entry is the
# only record of it until the rule can follow the chain across models.
allowlist_models = ["lighton_ocr"]

[rules.TRF002.explanation]
what_it_does = "Checks base_model_prefix is a non-empty, whitespace-free string literal."
why_bad = "Invalid prefixes can break weight-loading key mapping and base-model access."
diff = '''
 class AcmePreTrainedModel(PreTrainedModel):
-    base_model_prefix = ""
+    base_model_prefix = "model"
'''

[rules.TRF003]
description = "forward() should use @capture_output/@can_return_tuple, not manual return_dict branching."
default_enabled = false
allowlist_models = []

[rules.TRF003.explanation]
what_it_does = "Flags the old `if not return_dict: return (x,)` pattern in forward."
why_bad = "Manual return_dict branching is verbose and easy to get wrong. Let @capture_output or @can_return_tuple do it."
diff = '''
-def forward(self, x, return_dict=None):
-    if not return_dict:
-        return (x,)
-    return AcmeModelOutput(last_hidden_state=x)
+@can_return_tuple
+def forward(self, x):
+    return AcmeModelOutput(last_hidden_state=x)
'''

[rules.TRF004]
description = "Models must never override tie_weights. Use _tied_weights_keys instead."
default_enabled = true
# The four models that actually author a `tie_weights` override. `hubert`, `sew`, `unispeech`,
# `unispeech_sat` and `wavlm` were dropped: their modular files define no `tie_weights` at all, so the
# override exists only in their generated `modeling_*.py`, copied in from `Wav2Vec2ForCTC`. Fixing wav2vec2
# fixes all five, and listing them separately implied five independent violations instead of one root cause.
allowlist_models = ["data2vec", "sew_d", "wav2vec2", "wav2vec2_conformer"]

[rules.TRF004.explanation]
what_it_does = "Checks that no model class defines a tie_weights method."
why_bad = "Overriding tie_weights breaks loading, device_map computation and saving. Declare tied weights in the _tied_weights_keys class attribute."
diff = '''
-def tie_weights(self):
-    self.lm_head.weight = self.emb.weight
+class AcmeForCausalLM(AcmePreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
'''

[rules.TRF005]
description = "_no_split_modules, when defined, should be a list/tuple of non-empty strings."
default_enabled = true
# Empty: every previous entry (d_fine, deformable_detr, glm46v, lw_detr, pp_doclayout_v3, rt_detr,
# rt_detr_v2, voxtral, voxtral_realtime) is compliant under the current rule, which accepts an explicit
# `_no_split_modules = None` and an `AttributeError` sentinel in a modular file, and treats raw-string
# elements as the string constants they are. Verified by running the rule directly against each model.
allowlist_models = []

[rules.TRF005.explanation]
what_it_does = "Checks the shape of _no_split_modules when present."
why_bad = "Malformed values can break device-map partitioning and sharding."
diff = '''
-_no_split_modules = [SomeLayerClass, ""]
+_no_split_modules = ["AcmeDecoderLayer", "AcmeAttention"]
'''

[rules.TRF006]
description = "Cache arguments declared in a forward signature should be used in its body."
default_enabled = true
# `chinese_clip` was dropped: it no longer fires, verified by running the rule directly against its source.
allowlist_models = ["evolla", "idefics2", "llama4"]

[rules.TRF006.explanation]
what_it_does = "Checks that cache arguments in a forward signature are used in the body."
why_bad = "An unused cache argument suggests incomplete caching support and inconsistent API behavior."
diff = '''
 def forward(self, x, past_key_values=None, use_cache=False):
+    if use_cache:
+        ...
     return x
'''

[rules.TRF007]
description = "self.post_init() should stay last in a PreTrainedModel __init__."
default_enabled = true
allowlist_models = ["distilbert", "lxmert", "mt5", "pix2struct", "pop2piano", "switch_transformers", "t5"]

[rules.TRF007.explanation]
what_it_does = "Checks for self attribute assignments after self.post_init() in __init__."
why_bad = "Mutating model structure after post_init bypasses its initialization and finalization work."
diff = '''
 def __init__(self, config):
     ...
-    self.post_init()
-    self.proj = nn.Linear(...)
+    self.proj = nn.Linear(...)
+    self.post_init()
'''

[rules.TRF008]
description = "add_start_docstrings on a model class should not be empty."
default_enabled = true

[rules.TRF008.explanation]
what_it_does = "Checks that add_start_docstrings on a model class gets a non-empty argument."
why_bad = "An empty argument produces unclear, low-quality generated API docs."
diff = '''
-@add_start_docstrings("")
+@add_start_docstrings("The Acme model.")
 class AcmeModel(AcmePreTrainedModel):
     ...
'''

[rules.TRF009]
description = "No file in a model directory may import implementation code from another model package."
default_enabled = true
# 0.1.5 extended this rule from modeling files to every model file, surfacing 16 findings. Six were fixed by
# routing the sub-config through CONFIG_MAPPING/AutoConfig as the rule prescribes (cohere_asr, esmfold2,
# gemma4_assistant, gemma4_unified_assistant, ovis2, sam3's config). The ten below are allowlisted, in three
# groups:
#
# Back-compat tokenizer shims: `bart`, `convbert`, `distilbert`, `fnet`, `mobilebert`, `squeezebert`. Each
# tokenization_*.py exists only to re-export or thinly subclass another model's tokenizer, so the import is
# the module's entire purpose -- `BartTokenizer = RobertaTokenizer` cannot be written without importing
# RobertaTokenizer. The alternatives are deleting a public name or copying a whole WordPiece/BPE
# implementation into six files. Same population as the TRF038 entries below.
#
# generation_*.py: `nemotron3_5_asr`, `nemotron_asr_streaming`, `gradium_phonon_tts`. They subclass another model's
# generation mixin (RNNT mixin and cache; Pocket TTS's flow sampling loop), which the rule says belongs in a modular file -- but the modular converter has no
# generation_*.py support (no generated one exists in the repo, and the converter never mentions the file
# kind), so the prescribed remedy is unavailable. Either the converter learns this file kind or the rule
# stops checking it; flagged upstream rather than worked around here.
#
# `sam3`: modeling_sam3.py instantiates CLIPTextModelWithProjection. AutoModel.from_config would give
# CLIPTextModel, which covers the forward pass (sam3 reads only last_hidden_state and applies its own
# text_projection), but the released checkpoints carry `text_encoder.text_projection.weight`
# (convert_sam3_to_hf.py:69), a key only the WithProjection class owns. Swapping the class would orphan it.
# `shieldgemma2`: processing_shieldgemma2.py subclasses Gemma3Processor. The converter does support
# processing files, so this one is genuinely fixable with a modular_shieldgemma2.py -- left as follow-up work
# rather than folded into this pass.
#
# `vibevoice`: added after this pass, when VibeVoice (#40546) landed on main. generation_vibevoice.py imports
# `VibeVoiceAcousticTokenizerConv1dPaddingCache` from its own sub-model package. vibevoice is a composite
# family split across three directories (vibevoice, vibevoice_acoustic_tokenizer, vibevoice_asr) and already
# routes its sub-*configs* through CONFIG_MAPPING as the rule prescribes; this last import is a runtime cache
# class, which CONFIG_MAPPING cannot supply. It also sits in a generation_*.py file, so it hits the same
# missing converter support as the nemotron entries above.
allowlist_models = ["bart", "convbert", "distilbert", "dpr", "fnet", "gradium_phonon_tts", "mobilebert", "nemotron3_5_asr", "nemotron_asr_streaming", "sam3", "sam3_video", "shieldgemma2", "squeezebert", "vibevoice", "vision_text_dual_encoder"]

[rules.TRF009.explanation]
what_it_does = "Flags imports into another model's package from a model's shipped implementation files: modeling_*, configuration_*, processing_*, image_processing_*, video_processing_*, feature_extraction_*, tokenization_* and generation_*.py. Three forms count: package path (`from transformers.models.other.modeling_other import X`), relative path (`from ..other.configuration_other import X`) and public API (`from transformers import OtherModel`) -- for the last the owning directory is inferred from the class name and confirmed against the classes it defines, so an unresolvable name is left alone. Out of scope: `modular_*.py` (building on another model is its purpose, and the converter flattens those imports away), `convert_*.py`, `__init__.py` and files under `auto`. Allowed targets: the model's own directory, `auto` and `timm_wrapper`."
why_bad = "One model, one definition: model A's behavior lives in its own files, and a change to model B must not silently change model A -- for every file kind, not just modeling. Reach sub-configs through AutoConfig and CONFIG_MAPPING, and use a modular file to build on another model's code."
diff = '''
-from transformers.models.llama.modeling_llama import LlamaAttention
-from transformers import CLIPTextModelWithProjection
+# Keep implementation local to this model's own files.
+# To build on another model, write modular_acme.py; to reuse a snippet,
+# copy it with a # Copied from comment.
'''

[rules.TRF010]
description = "Direct config definitions must use @strict(accept_kwargs=True)."
default_enabled = true
# Empty: `nemotron_h` and `vibevoice_asr` both carry `@strict(accept_kwargs=True)` now.
allowlist_models = []

[rules.TRF010.explanation]
what_it_does = "Checks that direct PreTrainedConfig/PretrainedConfig subclasses in configuration_*.py and modular_*.py carry @strict(accept_kwargs=True)."
why_bad = "Without it a new config misses the repo's runtime type-validation contract and drifts from the dataclass-based config standard."
diff = '''
+@strict(accept_kwargs=True)
 class AcmeConfig(PreTrainedConfig):
     ...
'''

[rules.TRF011]
description = "forward() must not read non-nn.Module attributes off submodules: pipeline parallelism may replace them with Identity."
default_enabled = true
allowlist_models = []

[rules.TRF011.explanation]
what_it_does = "In forward() of PreTrainedModel subclasses, flags submodule attribute accesses torch.nn.Identity would not have: on loop variables over self.layers, and self.<submodule>.<attr> where <attr> is not a standard nn.Module attribute."
why_bad = "Pipeline parallelism may replace any submodule with torch.nn.Identity, so reading a custom attribute (e.g. decoder_layer.attention_type) off it raises AttributeError at runtime. Read per-layer metadata from self.config."
diff = '''
 def forward(self, ...):
-    for decoder_layer in self.layers:
+    for i, decoder_layer in enumerate(self.layers):
         hidden_states = decoder_layer(
             hidden_states,
-            attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+            attention_mask=causal_mask_mapping[self.config.layer_types[i]],
         )
'''

[rules.TRF012]
description = "_init_weights must use init primitives, not in-place operations on module weights."
default_enabled = true
allowlist_models = []

[rules.TRF012.explanation]
what_it_does = "Flags in-place ops (.normal_(), .zero_(), ...) on module weights inside _init_weights."
why_bad = "Parameters carry internal flags tracking whether they still need re-initialization; in-place ops bypass them. Use the `init` primitives."
diff = '''
+from transformers import initialization as init
+
 def _init_weights(self, module):
-    module.weight.normal_(mean=0.0, std=0.02)
+    init.normal_(module.weight, mean=0.0, std=0.02)
'''

[rules.TRF013]
description = "PreTrainedModel __init__ must call self.post_init()."
default_enabled = true
allowlist_models = []

[rules.TRF013.explanation]
what_it_does = "Checks that every PreTrainedModel subclass defining __init__ calls self.post_init(). In modular files super().__init__() counts, since it propagates the parent's post_init."
why_bad = "post_init does essential finalization (weight init, gradient checkpointing setup, ...); skipping it causes subtle runtime bugs."
diff = '''
 class AcmeModel(AcmePreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.layers = nn.ModuleList(...)
+        self.post_init()
'''

[rules.TRF014]
description = "`trust_remote_code` should never be used in native model integrations."
default_enabled = true
# The `auto` package is the loader layer that *implements* the `trust_remote_code`
# feature, forwarding the user-supplied flag to `*.from_pretrained`. It is not a native
# model integration, so it is exempt from this rule.
# The `auto` package is the loader layer that *implements* the `trust_remote_code`
# feature, forwarding the user-supplied flag to `*.from_pretrained`. It is not a native
# model integration, so it is exempt from this rule.
# The `auto` package is the loader layer that *implements* the `trust_remote_code`
# feature, forwarding the user-supplied flag to `*.from_pretrained`. It is not a native
# model integration, so it is exempt from this rule.
allowlist_models = ["auto"]

[rules.TRF014.explanation]
what_it_does = "Flags `trust_remote_code` used or passed (e.g. as a kwarg) in native model integration files."
why_bad = "`trust_remote_code` loads arbitrary code, including binaries -- a power feature for users, not something a native integration may depend on, since remote code cannot be reviewed or maintained in transformers."
diff = '''
 class AcmeModel(AcmePreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
-        self.model = AutoModel.from_pretrained(..., trust_remote_code=True)
+        self.model = AutoModel.from_pretrained(...)
'''

[rules.TRF015]
description = "Models with non-empty _tied_weights_keys must have tie_word_embeddings in their Config."
default_enabled = true
allowlist_models = []

[rules.TRF015.explanation]
what_it_does = "When a PreTrainedModel subclass sets a non-empty _tied_weights_keys, checks the companion configuration file for a tie_word_embeddings field."
why_bad = "Without it users cannot control weight tying: the model ties unconditionally, breaking serialization round-trips and fine-tuning with untied heads."
diff = '''
 # configuration_foo.py
 @strict(accept_kwargs=True)
 class FooConfig(PreTrainedConfig):
     hidden_size: int = 768
+    tie_word_embeddings: bool = True
'''

[rules.TRF016]
description = "do_* flags declared on a processor class must be referenced by overridden preprocess/_preprocess."
default_enabled = true
# Allowlist inherited unchanged from the base pin (0.1.2), so this PR changes no TRF001-TRF019 behavior.
# Both entries are the same shape: `do_convert_rgb` is declared on the video processor
# (`video_processing_kimi_k25.py:98`, `video_processing_minimax_m3_vl.py:45`) but never read by the
# `_preprocess` override, so only the model owner can say whether the flag is dead or the override lost a
# step. Left to the follow-ups with the two owners.
allowlist_models = ["kimi_k25", "minimax_m3_vl", "ernie4_5_vl_moe", "glm4v", "instructblipvideo", "internvl", "qwen2_vl", "smolvlm", "gemma4", "minicpmv4_6"]

[rules.TRF016.explanation]
what_it_does = "When an image_processing_*.py or video_processing_*.py class declares boolean do_* attributes (do_resize, do_rescale, do_normalize, do_convert_rgb, ...) and overrides preprocess() or _preprocess(), checks each flag is still consumed there: referenced directly, delegated via super().preprocess/_preprocess(..., **kwargs), or -- image processors only -- forwarded through _preprocess_image_like_inputs/_prepare_image_like_inputs. do_sample_frames is exempt: the base preprocess() consumes it before _preprocess() runs."
why_bad = "A do_X the override never references is dead: setting do_X=False has no effect, the operation runs anyway, and per-call overrides silently break."
diff = '''
 class AcmeVideoProcessor(BaseVideoProcessor):
     do_resize = True
     do_normalize = True

     def _preprocess(
         self,
         videos,
+        do_resize: bool,
+        do_normalize: bool,
         size,
         image_mean,
         image_std,
         **kwargs,
     ):
         for video in videos:
-            video = self.resize(video, size=size)
-            video = self.normalize(video, image_mean, image_std)
+            if do_resize:
+                video = self.resize(video, size=size)
+            if do_normalize:
+                video = self.normalize(video, image_mean, image_std)
'''

[rules.TRF017]
description = "@auto_docstring must be placed above @dataclass on output classes."
default_enabled = true
allowlist_models = []

[rules.TRF017.explanation]
what_it_does = "On classes carrying both @auto_docstring and @dataclass, checks @auto_docstring comes first."
why_bad = "Decorators apply bottom-up, so @dataclass on top runs @auto_docstring first, on a class with no synthesized __init__ yet: it then modifies the parent's __init__.__doc__ instead of the subclass's."
diff = '''
-@dataclass
 @auto_docstring(
     custom_intro="""
     Output type of [`AcmeForPreTraining`].
     """
 )
+@dataclass
 class AcmeForPreTrainingOutput(ModelOutput):
     ...
'''

[rules.TRF018]
description = "_init_weights overrides should call super()._init_weights(module), except modular-file sentinels."
default_enabled = true
# `radio` is the only entry left, and the one real case: it never calls the base implementation. The other
# 15 spelled the call `PreTrainedModel._init_weights(module)` rather than `super()._init_weights(module)` --
# the form a modular file needs, since `super()` there resolves to the inherited model, not to
# `PreTrainedModel`. The rule used to recognize only the `super()` spelling; that was reported upstream as
# huggingface/transformers-mlinter#27 and fixed there, so those 15 no longer fire and were dropped.
allowlist_models = ["radio"]

[rules.TRF018.explanation]
what_it_does = "Checks that every PreTrainedModel subclass overriding `_init_weights(self, module, ...)` chains up via `super()._init_weights(...)`. Modular files also accept the sentinels `PreTrainedModel._init_weights(self, module)`, `PreTrainedModel._init_weights(module)` and `raise AttributeError(...)`. For a deliberate full override, suppress with `# trf-ignore: TRF018` above the method."
why_bad = "The base `_init_weights` covers the standard module types (Linear, Embedding, LayerNorm, RotaryEmbedding, ...). Skipping it leaves every submodule the override misses uninitialized -- which passes tests and surfaces much later as a subtle weight-init bug (cf. https://github.com/huggingface/transformers/pull/45597)."
diff = '''
 from ... import initialization as init

 def _init_weights(self, module):
+    super()._init_weights(module)
     if isinstance(module, AcmeCustomLayer):
-        module.gate.data.zero_()
+        init.zeros_(module.gate)
'''

[rules.TRF019]
description = "A processor TypedDict must not define `_defaults`; ship them on the hub in `processor_config.json`."
default_enabled = true
# Allowlist inherited unchanged from the base pin (0.1.2). Three entries still fire -- `hunyuan_vl`,
# `kosmos2`, `qwen3_asr` -- and each sets `_defaults` on its processor kwargs class; dropping it means
# publishing those defaults in `processor_config.json` on the Hub repos first, which is on the checkpoint
# owners. `blip_2`, `grounding_dino` and `omdet_turbo` are covered by `cutoff_date` and no longer fire.
allowlist_models = ["blip_2", "grounding_dino", "hunyuan_vl", "kosmos2", "omdet_turbo", "qwen3_asr"]
cutoff_date = "2026-06-20"

[rules.TRF019.explanation]
what_it_does = "Flags a non-empty `_defaults` on `*ProcessorKwargs` TypedDicts in `processing_*.py`. Models released before the cutoff date are grandfathered."
why_bad = "Hardcoded `_defaults` scatter processor configuration across Python source, are awkward to override from config, and bloat the code. In `processor_config.json` on the hub they travel with the checkpoint and can be updated without a code change."
diff = '''
 class Gemma4ProcessorKwargs(ProcessingKwargs, total=False):
-    _defaults = {
-        "text_kwargs": {"padding": False},
-        "images_kwargs": {"return_tensors": "pt"},
-    }
     images_kwargs: Gemma4ImageProcessorKwargs
'''

[rules.TRF020]
description = "MLA models must isolate the KV LoRA expansion (kv_b_proj) in a dedicated method that forward() calls."
default_enabled = true
# Empty since mlinter#27 landed. The one entry, `axk2`, was a false positive: `AXK2Attention.forward` does
# call `self.expand_kv(k_pass, k_rot)`, and `expand_kv` is inherited from `DeepseekV32Attention`. The rule
# used to resolve the base chain only inside the linted file, so the modular file looked like it was missing
# the method; it now follows the chain across models and the finding is gone.
allowlist_models = []

[rules.TRF020.explanation]
what_it_does = "In model directories whose configuration declares `kv_lora_rank` (Multi-head Latent Attention), checks the attention class owning the KV LoRA expansion projection (`kv_b_proj`, or any `nn.Linear(config.kv_lora_rank, ...)`): the expansion must live in a dedicated method (e.g. `expand_kv`) that `forward()` calls, not inline. In modular files a method inherited from an imported base counts."
why_bad = "External backends (vLLM/SGLang) override the expansion to consume the compressed KV cache directly. Inlined in `forward()` there is nothing to override, so the backend must materialize the full cache -- losing the memory savings MLA exists for."
diff = '''
+    def expand_kv(self, k_nope, k_pe):
+        key_shape = (*k_nope.shape[:-1], -1, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope = self.kv_b_proj(k_nope).view(key_shape).transpose(1, 2)
+        k_nope, value_states = torch.split(k_nope, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+        k_pe = k_pe.expand(*k_nope.shape[:-1], -1)
+        key_states = torch.cat((k_nope, k_pe), dim=-1)
+        return key_states, value_states
+
     def forward(self, hidden_states, ...):
         ...
-        k_nope = self.kv_b_proj(k_pass).view(key_shape).transpose(1, 2)
-        k_nope, value_states = torch.split(k_nope, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
-        k_pe = k_rot.expand(*k_nope.shape[:-1], -1)
-        key_states = torch.cat((k_nope, k_pe), dim=-1)
+        key_states, value_states = self.expand_kv(k_pass, k_rot)
'''

[rules.TRF021]
description = "Scalar tensors must be filled on-device with torch.full((), ...) instead of torch.tensor(..., device=...)."
default_enabled = true
allowlist_models = []

[rules.TRF021.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags `torch.tensor(<value>, ..., device=<non-cpu>)` where `<value>` provably resolves to a Python scalar -- from numeric literals and arithmetic on them, torch.finfo/iinfo fields, scalar-returning builtins and math.* calls, locals bound exactly once, self.<attr> assigned in the class body, and config fields annotated int/float/bool in the companion configuration file (following attribute_map). Anything that may also be a sequence (`eos_token_id: int | list[int] | None`) or that cannot be resolved is left alone. __init__, _init_weights, __post_init__ and post_init are exempt: they never run inside a capture region."
why_bad = "torch.tensor(<python scalar>, device=<accelerator>) materialises the value on the host then copies it to the device; CUDA graph capture forbids that copy, so the model cannot be captured. torch.full((), <value>, dtype=..., device=...) fills the same 0-d tensor on-device with a capturable kernel and no synchronisation."
diff = '''
 def get_placeholder_mask(self, input_ids, inputs_embeds):
     special_image_mask = (
         inputs_embeds
         == self.get_input_embeddings()(
-            torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+            torch.full((), self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
         )
     ).all(-1)
'''

[rules.TRF022]
description = "_no_split_modules entries must name module classes that exist in the model."
default_enabled = true
allowlist_models = []

[rules.TRF022.explanation]
what_it_does = "Checks that every string in a `_no_split_modules` list in modeling_*.py or modular_*.py names a class defined in that file, imported into it, or defined by a sibling module of the same model directory. Complements TRF005, which only checks the value's shape."
why_bad = "`device_map` matches these strings against `module.__class__.__name__` at runtime, so a stale or misspelled name matches nothing and is silently ignored -- the module it should keep together can still be split across devices. Delete entries naming another model's classes rather than correcting them: `post_init` already collects `_no_split_modules` from child submodels."
diff = '''
 class VideoLlavaPreTrainedModel(PreTrainedModel):
-    _no_split_modules = ["VideoLlavaVisionAttention"]
'''

[rules.TRF023]
description = "Config fields must use canonical dimension names (hidden_size, intermediate_size, num_attention_heads, ...)."
default_enabled = true
# openai -> openai-gpt.md and kosmos2 -> kosmos-2.md cannot be derived from the directory name, so the
# cutoff cannot grandfather them. qwen3_asr mirrors Whisper's encoder config, where d_model is public API.
# openai -> openai-gpt.md and kosmos2 -> kosmos-2.md cannot be derived from the directory name, so the
# cutoff cannot grandfather them. qwen3_asr mirrors Whisper's encoder config, where d_model is public API.
# openai -> openai-gpt.md and kosmos2 -> kosmos-2.md cannot be derived from the directory name, so the
# cutoff cannot grandfather them. qwen3_asr mirrors Whisper's encoder config, where d_model is public API.
allowlist_models = ["kosmos2", "openai", "qwen3_asr"]
cutoff_date = "2026-06-20"

[rules.TRF023.explanation]
what_it_does = "In configuration_*.py and modular_*.py, flags `*Config` fields named after an upstream paper's abbreviation instead of the canonical name: d_model/n_embd -> hidden_size, d_ff/d_inner/ffn_dim/ffn_hidden_size/expansion_ratio -> intermediate_size, d_head -> head_dim, n_head/n_heads -> num_attention_heads, n_layer/n_layers/num_blocks -> num_hidden_layers. Fields are read from the class body and from __init__/__post_init__ assignments and defaults. Ambiguous but idiomatic names (num_heads, num_layers, embed_dim, mlp_ratio) are not flagged, and models added before cutoff_date keep theirs."
why_bad = "Everything that reads a model's shape -- device_map planning, tensor/pipeline parallel plans, quantization, PEFT, attention-backend selection, `attribute_map` consumers -- looks up the canonical names, so a config spelling the same quantity `d_model` silently opts out of all of them. Map the checkpoint's own spelling in the conversion script."
diff = '''
 @strict(accept_kwargs=True)
 class AcmeConfig(PreTrainedConfig):
-    d_model: int = 1024
-    d_ff: int = 4096
-    n_heads: int = 16
-    n_layers: int = 24
+    hidden_size: int = 1024
+    intermediate_size: int = 4096
+    num_attention_heads: int = 16
+    num_hidden_layers: int = 24
'''

[rules.TRF024]
description = "Layer dimensions must come from the config, not from an integer literal in the modeling file."
default_enabled = true
allowlist_models = []
cutoff_date = "2026-06-20"

[rules.TRF024.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags an integer literal greater than 8 in a dimension argument of a torch.nn constructor (Linear, Embedding, LayerNorm, RMSNorm, GroupNorm, BatchNorm*, InstanceNorm*, Conv*d, ConvTranspose*d, Bilinear, MultiheadAttention), positional or by keyword (in_features, out_features, in_channels, out_channels, num_embeddings, embedding_dim, embed_dim, normalized_shape, num_channels, hidden_size). Operator-shape arguments (kernel_size, stride, padding, num_groups) are ignored; literals up to 8 keep scalar heads, binary classifiers and RGB channel counts clean. Models added before cutoff_date are exempt."
why_bad = "A hardcoded width pins the module to one checkpoint size: the same architecture at another scale loads with a shape mismatch, and `from_pretrained` cannot say which value is wrong because no config field points at it. It also splits the source of truth, so editing the config no longer changes the model that gets built."
diff = '''
 class AcmeAtomEmbedding(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.proj = nn.Linear(768, 3072, bias=False)
-        self.norm = nn.LayerNorm(3072)
+        self.proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.norm = nn.LayerNorm(config.intermediate_size)
'''

[rules.TRF025]
description = "Attention masks must be built once in the model, not rebuilt inside a layer or attention module."
default_enabled = true
allowlist_models = []
cutoff_date = "2026-06-20"

[rules.TRF025.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags calls to a mask factory (the masking_utils entry points create_causal_mask, create_bidirectional_mask, create_sliding_window_causal_mask, create_chunked_causal_mask, create_masks_for_generate, and any create_*_mask helper) inside a class that does not inherit from PreTrainedModel -- plain nn.Module blocks such as layers, attention modules and encoders."
why_bad = "Mask construction is O(sequence length squared) work that does not vary per layer, so building it in the layer pays that cost once per layer, and each layer then owns its own mask, so the attention backends can no longer be handed a single prepared one. Build it once in the model and pass it down."
diff = '''
 class AcmeLayer(nn.Module):
     def forward(self, hidden_states, attention_mask=None, **kwargs):
-        attention_mask = create_causal_mask(
-            config=self.config, input_embeds=hidden_states, attention_mask=attention_mask, ...
-        )
         return self.self_attn(hidden_states, attention_mask, **kwargs)

 class AcmeModel(AcmePreTrainedModel):
     def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        causal_mask = create_causal_mask(
+            config=self.config, input_embeds=inputs_embeds, attention_mask=attention_mask, ...
+        )
         for layer in self.layers:
-            hidden_states = layer(hidden_states, attention_mask, **kwargs)
+            hidden_states = layer(hidden_states, causal_mask, **kwargs)
'''

[rules.TRF026]
description = "A module whose forward only delegates to its single submodule adds nothing; inline it."
default_enabled = true
allowlist_models = []
cutoff_date = "2026-06-20"

[rules.TRF026.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags a non-PreTrainedModel class that defines only __init__ and forward, assigns exactly one self.<attr> in __init__, and whose forward body is exactly `return self.<attr>(...)` for it (a leading docstring is ignored). Any other method, extra attribute, statement before the return, or `super()` call in forward means the class does work of its own."
why_bad = "The wrapper adds a level to every weight name, to _no_split_modules, parallelism plans and every conversion mapping, while computing nothing -- and readers have to open one more class to find that out. PreTrainedModel subclasses are exempt: they exist for from_pretrained and the auto classes even when forward only delegates."
diff = '''
-class AcmeAtomTransformer(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.encoder = AcmeEncoder(config)
-
-    def forward(self, hidden_states, **kwargs):
-        return self.encoder(hidden_states, **kwargs)
-
 class AcmeModel(AcmePreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
-        self.atom_transformer = AcmeAtomTransformer(config)
+        self.encoder = AcmeEncoder(config)
'''

[rules.TRF027]
description = "Model files must raise explicit errors instead of using bare `assert`."
default_enabled = true
allowlist_models = []
cutoff_date = "2026-06-20"

[rules.TRF027.explanation]
what_it_does = "Flags any `assert` statement in modeling_*.py, modular_*.py and configuration_*.py."
why_bad = "`python -O` strips asserts, so a shape or config check written as one silently disappears in optimised runs. An assert also gives a bare AssertionError, where a ValueError can name the offending value and say what to do about it."
diff = '''
 def forward(self, hidden_states):
-    assert hidden_states.dim() == 3
+    if hidden_states.dim() != 3:
+        raise ValueError(f"Expected a 3D tensor, got shape {tuple(hidden_states.shape)}.")
'''

[rules.TRF028]
description = "Model files must carry a complete license header."
default_enabled = true
allowlist_models = []
cutoff_date = "2026-06-20"

[rules.TRF028.explanation]
what_it_does = "Checks the first 25 lines of modeling_*.py, modular_*.py, configuration_*.py, processing_*.py, image_processing_*.py and video_processing_*.py for a `Licensed under the <name> License` line followed by every clause of the standard warranty paragraph, from `You may obtain a copy of the License at` through `limitations under the License.`. Lines are flattened and lowercased first, so wrapping and comment style do not matter. The license name and copyright line are not checked: they vary per model."
why_bad = "A file shipped without the header leaves its provenance ambiguous, and adding it later means touching an already-released file. Matching only `Apache License` would let through the defects that actually occur: a paragraph truncated mid-way, a header that stops after the URL, or one mangled by a bad search-and-replace."
diff = '''
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# ...
 """PyTorch Acme model."""
'''

[rules.TRF029]
description = "A module taking `config` must not also take arguments that live on the config."
default_enabled = true
allowlist_models = ["kosmos2"]
cutoff_date = "2026-06-20"

[rules.TRF029.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags an `__init__` taking `config` alongside an argument whose name is unambiguously a config field (hidden_size, num_attention_heads, intermediate_size, head_dim, num_hidden_layers, embed_dim, dropout, eps, patch_size, rope_theta, ...). A parameter optional with a `None` default is exempt: that is an override, not a second source of truth, and it is how one MLP class serves both the dense and the expert width of a MoE model. A hardcoded default such as `hidden_size: int = 1024` is not -- it wins over the config whenever the caller passes nothing. kosmos2 is allowlisted: its doc page is not derivable from the directory name, so the cutoff cannot grandfather it."
why_bad = "The same number now has two sources of truth and the caller picks the winner, so editing the config no longer changes the model that gets built. It also pushes architecture knowledge out to every call site, where it does not belong."
diff = '''
 class AcmeAttention(nn.Module):
-    def __init__(self, config, embed_dim, num_heads, dropout):
+    def __init__(self, config, layer_idx=None):
         super().__init__()
-        self.embed_dim = embed_dim
-        self.num_heads = num_heads
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads

 class AcmeMLP(nn.Module):
     # an optional override is fine: omitting it reads the config
     def __init__(self, config, intermediate_size=None):
         super().__init__()
         self.intermediate_size = intermediate_size or config.intermediate_size
'''

[rules.TRF030]
description = "Reaching more than two levels into the config means the module was handed the wrong config."
default_enabled = true
allowlist_models = []
cutoff_date = "2026-06-20"

[rules.TRF030.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags attribute chains rooted at `config` or `self.config` three or more levels deep. `config.hidden_size` (one hop) and `config.text_config.hidden_size` (two hops, the normal sub-config access) are fine. One violation per line."
why_bad = "A module that walks `config.diffusion_config.atom_encoder_config.hidden_size` is coupled to the whole config hierarchy rather than its own slice, so it cannot be reused, tested or given a different sub-config. Pass the relevant sub-config down and the chain collapses to one hop."
diff = '''
 class AcmeAtomEncoder(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.norm = AcmeLayerNorm(config.diffusion_config.atom_encoder_config.hidden_size)
+        self.norm = AcmeLayerNorm(config.hidden_size)
'''

[rules.TRF031]
description = "Output-like dataclasses in modeling files must inherit ModelOutput."
default_enabled = true
allowlist_models = []
cutoff_date = "2026-06-20"

[rules.TRF031.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags a top-level `@dataclass` whose bases include nothing ending in `Output`, unless it has two or more mandatory fields -- those are internal argument bundles, which `ModelOutput` rejects at runtime."
why_bad = "A plain output dataclass does not index like a tuple, does not survive `return_dict=False`, and is invisible to @auto_docstring, so its fields never reach the generated API docs. ModelOutput gets all three for free."
diff = '''
 @auto_docstring
 @dataclass
-class AcmeStructureOutput:
+class AcmeStructureOutput(ModelOutput):
     positions: torch.Tensor
     confidence: Optional[torch.Tensor] = None
'''

[rules.TRF032]
description = "Masked positions must be filled with torch.finfo(dtype).min, not a magic negative number."
default_enabled = true
allowlist_models = []
cutoff_date = "2026-06-20"

[rules.TRF032.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags masked_fill, masked_fill_, full, full_like and new_full called with a negated numeric literal of magnitude 1e3 or more."
why_bad = "A hardcoded -1e9 overflows to -inf in float16 and is nowhere near the smallest value in float32, so the same mask behaves differently per dtype and can produce NaNs after softmax. `torch.finfo(dtype).min` is the smallest representable value in whatever dtype is actually running."
diff = '''
-attention_scores = attention_scores.masked_fill(~mask, -1e9)
+attention_scores = attention_scores.masked_fill(~mask, torch.finfo(attention_scores.dtype).min)
'''

[rules.TRF033]
description = "Hyperparameters must be set on the config, not mutated through a set_* method."
default_enabled = true
allowlist_models = []
cutoff_date = "2026-06-20"

[rules.TRF033.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags methods whose name starts with `set_`, except the PreTrainedModel contract methods set_input_embeddings, set_output_embeddings, set_decoder, set_encoder, set_attn_implementation and set_default_language."
why_bad = "A setter makes behaviour depend on call order: the value is not in the config, so it is not saved, not restored by from_pretrained, and not visible to device-map or parallelism planning. Users have to know to call it, and forgetting is silent."
diff = '''
 class AcmeTriangleAttention(nn.Module):
-    def set_chunk_size(self, chunk_size):
-        self.chunk_size = chunk_size
+    def __init__(self, config):
+        super().__init__()
+        self.chunk_size = config.chunk_size
'''

[rules.TRF034]
description = "Layer classes held in an nn.ModuleList must subclass GradientCheckpointingLayer."
default_enabled = true
# `x_clip`'s `PromptGeneratorLayer` is all that is left: an attention stack on a model that does enable
# gradient checkpointing, so it is a real finding and a follow-up for that model.
# The list was 17 models and got there in two steps. mlinter#27 fixed base-chain resolution across model
# files, which cleared the 12 modular layers that only looked unpatched because their generated
# `modeling_*.py` did subclass `GradientCheckpointingLayer` (cohere_compass, cosmos3_edge, glm5_next,
# granite_swa, granitemoe_swa, hunyuan_vl, kimi_k25, muse_glimmer, muse_glimmer_assistant, radio, step3p7,
# tipsv2). mlinter#65 then scoped the rule to the checkpointable trunk, which cleared the remaining four:
# `dinov3_convnext`, `openai` and `xcodec2` never set `supports_gradient_checkpointing`, so nothing could be
# skipped, and `tipsv2_dpt`'s three entries are DPT head layers rather than a token-mixing stack.
allowlist_models = ["x_clip"]
cutoff_date = "2026-06-20"

[rules.TRF034.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags a locally-defined `*Layer`/`*Block` class instantiated in an `nn.ModuleList(...)` that does not reach `GradientCheckpointingLayer` through its base chain; modular files follow relative imports into sibling models, and unresolved chains are inconclusive. Out of scope: a model that never sets `supports_gradient_checkpointing = True`, which raises from `gradient_checkpointing_enable()` instead of skipping a layer; a layer holding `nn.BatchNorm*`/`nn.InstanceNorm*`, whose statistics would be recomputed twice; and a stack that is not the model's token mixer, shown by an attention, modulation, mixer or SSM module assigned as `self.x = Y(...)`."
why_bad = "`gradient_checkpointing_enable()` wraps a layer only if it is a GradientCheckpointingLayer. A plain nn.Module on the trunk is skipped silently, so training looks checkpointed while still allocating full activations, and the OOM surfaces far from the cause. Elsewhere -- a conv backbone, a decode head -- the trade is the author's call, so the rule stays out."
diff = '''
-class AcmeDecoderLayer(nn.Module):
+class AcmeDecoderLayer(GradientCheckpointingLayer):
     def __init__(self, config, layer_idx):
         super().__init__()
'''

[rules.TRF035]
description = "Model files must not silence the linter with `# noqa`."
default_enabled = true
# Empty since the rule learned to accept F401/F821/F822 in modular files, which is what all five previously
# allowlisted models (hunyuan_vl, nemotron3_5_asr, step3p7, zaya, glm5_next) were silencing: a re-exported
# `repeat_kv`, a converter-steering import, a quoted forward reference to another model's cache class, and a
# `del` of a name that only exists in the inherited parent body. That was flagged upstream as converter
# mechanics rather than model bugs, and the fix landed there, so no model needs an exemption any more.
allowlist_models = ["hy_v4"]
cutoff_date = "2026-06-20"

[rules.TRF035.explanation]
what_it_does = "Flags `# noqa` comments, with or without codes, in modeling_*.py, modular_*.py and configuration_*.py. In a modular file `F401`, `F821` and `F822` are accepted: it is a generation source that deliberately does not define every name it uses, so ruff's undefined-name family fires on correct code -- `__all__` entries the converter fills in, classes living in the parent model, imports kept to be re-exported. A `# noqa` naming only those codes is skipped; one naming anything else is reported on the codes that are left. A bare `# noqa` is always reported; in a modular file the message asks for the code rather than a rewrite."
why_bad = "Model files are ordinary code held to the repo's lint rules, so a suppression means the underlying issue was left in place -- and a bare `# noqa` also hides every future violation on that line. Fix the code instead."
diff = '''
-from ...modeling_utils import PreTrainedModel  # noqa: F401
+from ...modeling_utils import PreTrainedModel
'''

[rules.TRF036]
description = "nn.Sequential hides the forward flow; declare the submodules explicitly."
default_enabled = true
allowlist_models = ["x_clip"]
cutoff_date = "2026-06-20"

[rules.TRF036.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags any `nn.Sequential(...)` construction."
why_bad = "Sequential names its children by position, so weights land at `mlp.0.weight` and `mlp.2.weight`: the conversion mapping, _tied_weights_keys and every parallelism plan then reference indices, and inserting a layer renames everything after it. It also hides the forward, so the dtype casts and residuals between steps are invisible where they happen."
diff = '''
-        self.mlp = nn.Sequential(
-            nn.Linear(config.hidden_size, config.intermediate_size),
-            nn.GELU(),
-            nn.Linear(config.intermediate_size, config.hidden_size),
-        )
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.act = ACT2FN[config.hidden_act]
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
'''

[rules.TRF037]
description = "torch.einsum is hard to read; prefer explicit matmul/transpose operations."
default_enabled = false
allowlist_models = ["x_clip"]
cutoff_date = "2026-06-20"

[rules.TRF037.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags calls to einsum, reporting the equation when it is a literal. Disabled by default: einsum is occasionally the clearest way to express a contraction, so this is opt-in rather than a hard convention."
why_bad = "An einsum equation encodes the shapes in a notation the reader has to decode, and a dynamically built equation hides which contraction runs at all. Expand it into explicit matmul/transpose operations."
diff = '''
-        pair_bias = torch.einsum("bqhc,bkhc->bhqk", query_states, key_states)
+        pair_bias = query_states.permute(0, 2, 1, 3) @ key_states.permute(0, 2, 3, 1)
'''

[rules.TRF038]
description = "Every modeling-family source file needs a matching test file under tests/models/."
default_enabled = true
# 49 findings over 38 models. Verified during review: only `nougat` and `glmga` genuinely have no test.
# For the rest the rule pairs each modeling-family source file with its own test file, so a family that
# has a processor / feature extractor but no dedicated test for it is flagged even though the model itself
# is tested. Agreed in review: enforce for new models, backfill the old families opportunistically, rather
# than block this PR on ~30 test files. `sam2` no longer fires and was dropped.
#
# `bart`, `mobilebert` and `squeezebert` are a different case, added when 0.1.5 extended this rule to
# tokenizer files. Their `tokenization_*.py` define no classes at all -- each is a two-line back-compat alias
# (`BartTokenizer = RobertaTokenizer`, `MobileBertTokenizer = BertTokenizer`, `SqueezeBertTokenizer =
# BertTokenizer`) left by "rm slow tokenizers" (#40936), which deleted their test files on purpose because the
# slow tokenizers they exercised no longer exist. A test file here could only assert alias identity, and the
# aliased tokenizers are already covered by tests/models/roberta and tests/models/bert. The better fix is
# upstream -- TRF038 should skip a source file that defines no classes -- so this is a placeholder, not a
# permanent exemption. The other tokenizer findings 0.1.5 surfaced (bert_legacy, convbert, dpr, fnet, qwen3_5)
# got real test files instead, and `roberta` was fixed by deleting its dead `tokenization_roberta_old.py`.
allowlist_models = ["bart", "bros", "cohere_asr", "emu3", "exaone4_5", "gemma4", "glm46v", "glm_image", "glmga", "granite_speech", "granite_speech5", "higgs_audio_v2", "inkling", "kyutai_speech_to_text", "lasr", "minimax_m3_vl", "mobilebert", "moonshine_streaming", "nemotron_asr_streaming", "nemotron3_5_asr", "nougat", "paddleocr_vl", "pe_audio", "pe_video", "pe_audio_video", "pi0", "pp_formulanet", "rf_detr", "sam2_video", "sam3_tracker", "sam3_tracker_video", "sam3_video", "siglip", "siglip2", "squeezebert", "tvp", "video_llava", "vilt", "vjepa2", "voxtral", "voxtral_realtime", "x_clip", "glm5_next"]

[rules.TRF038.explanation]
what_it_does = "Checks that every modeling_*.py, processing_*.py, image_processing_*.py, video_processing_*.py, feature_extraction_*.py and tokenization_*.py file has a matching tests/models/<model>/test_*.py (modeling_acme.py -> tests/models/acme/test_modeling_acme.py). Exempt: configuration_*.py, covered by ConfigTester in the companion test_modeling_*.py, and tokenization_utils*.py, a helper covered through its tokenizer; tokenization_<name>_fast.py maps to its slow counterpart's test file. A modular_*.py can define several families at once, so its classes are routed by name suffix -- XxxModel/XxxPreTrainedModel/XxxFor<Task> modeling, XxxImageProcessor(Fast) image processing, XxxProcessor processing, XxxVideoProcessor video processing, XxxFeatureExtractor feature extraction, XxxTokenizer(Fast) tokenization, XxxConfig skipped -- and one violation is reported per missing test file. `# trf-ignore: TRF038` is not supported; use `allowlist_models` so an exemption is visible in review."
why_bad = "A source file with no test file has no regression coverage: a broken forward pass, a bad conversion mapping, or a tokenizer that drops a special token can land and stay broken indefinitely. A dummy config with random weights, or a small hand-written vocabulary, is enough to exercise any of them."
diff = '''
 src/transformers/models/acme/modeling_acme.py
+tests/models/acme/test_modeling_acme.py
 src/transformers/models/acme/tokenization_acme.py
+tests/models/acme/test_tokenization_acme.py
'''

[rules.TRF039]
description = "Imports guarded by is_*_available() must be removed once nothing in the file uses them."
default_enabled = true
allowlist_models = []

[rules.TRF039.explanation]
what_it_does = "Finds `if is_*_available(): import ...` blocks (including combinations such as `is_vision_available() and is_torch_available()`) and flags the import when the name is referenced nowhere else in the file, including inside string type hints and __all__."
why_bad = "ruff's unused-import check does not clean these up: the import is reachable, so the block looks fine on its own. Once a refactor no longer needs PIL.Image, torch, etc., the guarded import lingers as dead weight and a misleading signal about the file's real dependencies."
diff = '''
 if is_vision_available():
-    from PIL import Image
'''

[rules.TRF040]
description = "@can_return_tuple must not be combined with @capture_outputs."
default_enabled = true
allowlist_models = []

[rules.TRF040.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags methods decorated with both @capture_outputs and @can_return_tuple. Complements TRF003, which covers manual return_dict branching in forward()."
why_bad = "Both decorators pop return_dict, so only the outermost one sees the true value. @capture_outputs already handles the to_tuple conversion, making @can_return_tuple redundant."
diff = '''
-@can_return_tuple
 @merge_with_config_defaults
 @capture_outputs
 @auto_docstring
 def forward(self, x):
     return AcmeModelOutput(last_hidden_state=x)
'''

[rules.TRF041]
description = "A config-gated branch must carry a `# CODEPATH:` note saying which checkpoints diverge."
default_enabled = false
# 15 findings over these 9 models: dinov3_vit(4), inkling(4), axk2(1), cohere_compass(1), kimi_k25(1),
# kosmos2(1), muse_glimmer(1), tipsv2(1), x_clip(1). Every one needs its `# CODEPATH:` note to name the
# checkpoints that take each path, which only the model owner knows -- guessing would be worse than the
# missing note -- so they go out as per-model follow-ups. `mimo_v2_flash` no longer fires and was dropped.
# Down from 35 findings over 13 models: the rule now exempts guards (an `if` with no `else` that only raises
# or only logs) and a `DEFAULT_EXEMPT_ATTRIBUTES` list of framework plumbing, which cleared `openai`(11),
# `timm_backbone`(2), `timm_wrapper`(2), `vitpose_backbone`(1) and `glm5_next` outright and trimmed the rest.
# `embedding_gemma2`(4): the optional vision/audio tower branches it inherits from `gemma4`, which gate on
# `config.vision_config is not None` / `config.audio_config is not None`. Whether a tower is configured is
# composition, not checkpoint divergence, so a `# CODEPATH:` note would restate the condition.
allowlist_models = ["axk2", "cohere_compass", "dinov3_vit", "embedding_gemma2", "inkling", "kimi_k25", "kosmos2", "muse_glimmer", "pp_doclayout_v4", "tipsv2", "x_clip", "hy_v4"]
cutoff_date = "2026-06-20"

[rules.TRF041.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags every `if`/`elif` and conditional expression whose condition reads a `config.*` or `self.config.*` attribute without a `# CODEPATH:` comment, accepted on the branch line or in the contiguous comment block directly above it. Any config attribute counts, not just boolean flags. Exempt by structure: `X if X is not None else fallback`, where the tested field is itself one of the results -- `getattr(config, x, default)` spelled long, which cannot fork the graph; and a guard, an `if` with no `else` whose body only raises or only warns/logs, since one side aborts and nothing diverges past it. Merely mentioning None does not qualify: `config.vision_config is not None` still owes a note. Exempt by field: framework plumbing that gates no checkpoint divergence -- `problem_type` selecting a loss, `hidden_act` looking up an activation, `num_labels`, `use_cache`, `is_decoder`, the special token ids, the `summary_*` head settings; the full list is `DEFAULT_EXEMPT_ATTRIBUTES` in `mlinter/trf041.py`, extended per project by `ignored_attributes = [...]` on the rule table. A model exempts one of its own fields file-wide with a module-level `# trf-ignore: TRF041 config.scale_embedding, config.auxiliary_loss` at column 0 (`self.config.x`, `config.x` and `x` are the same field). It must name at least one field -- a bare `# trf-ignore: TRF041` keeps its per-line meaning -- and a condition is skipped only when every field it reads is exempt."
why_bad = "A config-gated branch is a second architecture in the same file, and the code cannot say whether both halves are still reachable -- which is how dead experimental branches survive for releases. The rule does not forbid the branch: like Rust's `// SAFETY:`, it asks for the checkpoints taking each side to be written down next to it. A branch nobody can name one for is a branch to delete."
diff = '''
+        # CODEPATH: ESMC-6B ships pre-normalised embeddings, the 300M/600M checkpoints do not.
         if config.use_embedding_norm:
             hidden_states = self.embedding_norm(hidden_states)

-        if config.msa_encoder_enabled:
-            hidden_states = self.msa_encoder(hidden_states)
+        # no released checkpoint sets msa_encoder_enabled -> branch removed
'''

[rules.TRF042]
description = "A tokenizer test must exercise the shared TokenizerTesterMixin suite."
default_enabled = true
allowlist_models = ["auto"]
cutoff_date = "2026-06-20"

[rules.TRF042.explanation]
what_it_does = "In tests/models/*/test_tokenization_*.py, checks the file defines a collected test class inheriting TokenizerTesterMixin. Only classes the runner collects count -- a `TestCase` base, or the `*Test` naming convention when the base is another model's test class -- so helper-only files are skipped, and a helper mixing in the suite does not satisfy the rule for a real test class. Inheritance is followed through bases in the same file and into another model's tokenizer test imported by name; an unresolvable base never counts. Reported on the first test class that does not run the suite. `auto` is allowlisted: test_tokenization_auto.py tests AutoTokenizer resolution, not a tokenizer."
why_bad = "TokenizerTesterMixin is where encode/decode round-tripping, padding and truncation, special-token handling, added-token persistence and save/load equivalence are actually checked. A test that only asserts a few hand-written token id lists passes while the tokenizer is broken in every one of them, and still looks tested in review."
diff = '''
-class AcmeTokenizationTest(unittest.TestCase):
+class AcmeTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
+    tokenizer_class = AcmeTokenizer
+    test_slow_tokenizer = True
'''

[rules.TRF043]
description = "Attention classes must not declare position_ids in their forward signature; it flows via **kwargs."
default_enabled = true
# This rule has no `cutoff_date`, so it fires across the pre-existing corpus: 21 attention classes, one
# site each. 14 are legacy `modeling_*.py` (chameleon, clvp, codegen, falcon, gpt_neox_japanese, gptj,
# idefics, idefics2, moshi, persimmon, prophetnet, qwen2_vl, recurrent_gemma, stablelm), 7 are modular
# (axk2, deepseek_v32, deepseek_v4, gemma4, glm_moe_dsa, ministral3, mistral4). Dropping `position_ids`
# from a public `forward` signature is a breaking change per model, so each is its own follow-up. Upstream:
# a `cutoff_date` would express the legacy half better than a list this long.
allowlist_models = ["axk2", "chameleon", "clvp", "codegen", "deepseek_v32", "deepseek_v4", "falcon", "gemma4", "glm_moe_dsa", "gpt_neox_japanese", "gptj", "idefics", "idefics2", "ministral3", "mistral4", "moshi", "persimmon", "prophetnet", "qwen2_vl", "recurrent_gemma", "stablelm", "hy_v4"]

[rules.TRF043.explanation]
what_it_does = "Flags a declared position_ids parameter in the forward signature of classes whose name ends in Attention."
why_bad = "position_ids is consumed downstream by flash-attention padding-free training and must flow through **kwargs. Naming it in the signature swallows it before the attention interface can read it; the llama standard passes position_embeddings plus **kwargs."
diff = '''
 class AcmeAttention(nn.Module):
     def forward(
         self,
         hidden_states,
         position_embeddings,
         attention_mask=None,
-        position_ids=None,
         **kwargs: Unpack[TransformersKwargs],
     ):
'''

[rules.TRF044]
description = "cache_position must not appear as a parameter in modeling code; it is removed framework surface."
default_enabled = true
allowlist_models = []

[rules.TRF044.explanation]
what_it_does = "Flags a parameter named cache_position on any function in modeling_*.py and modular_*.py."
why_bad = "cache_position was removed from all models in v5. Reintroducing it (usually copied from pre-v5 sources) threads a dead argument through every layer; the cache update is past_key_values.update(key_states, value_states, self.layer_idx), with no position threading."
diff = '''
 def forward(
     self,
     hidden_states,
     past_key_values=None,
-    cache_position=None,
     **kwargs,
 ):
-    key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_position)
+    key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
'''

[rules.TRF045]
description = "forward must not declare output_attentions/output_hidden_states/return_dict; decorators own them."
default_enabled = true
# kosmos2 and openai have doc pages (kosmos-2.md, openai-gpt.md) whose names cannot be derived from the
# model directory, and the timm wrappers have no contribution date line, so the cutoff cannot exempt them.
# kosmos2 and openai have doc pages (kosmos-2.md, openai-gpt.md) whose names cannot be derived from the
# model directory, and the timm wrappers have no contribution date line, so the cutoff cannot exempt them.
# kosmos2 and openai have doc pages (kosmos-2.md, openai-gpt.md) whose names cannot be derived from the
# model directory, and the timm wrappers have no contribution date line, so the cutoff cannot exempt them.
allowlist_models = ["kosmos2", "openai", "timm_backbone", "timm_wrapper"]
cutoff_date = "2026-06-20"

[rules.TRF045.explanation]
what_it_does = "Flags the legacy output_attentions, output_hidden_states and return_dict parameters in forward signatures in modeling_*.py and modular_*.py. Models contributed before the cutoff date are exempt."
why_bad = "The decorator stack owns output control: @capture_outputs resolves the output_* flags against the config and records tensors via _can_record_outputs, and @can_return_tuple handles return_dict. Declaring them in the signature reintroduces manual flag threading that drifts from the decorators."
diff = '''
+@capture_outputs
 def forward(
     self,
     input_ids,
-    output_attentions=None,
-    output_hidden_states=None,
-    return_dict=None,
     **kwargs: Unpack[TransformersKwargs],
 ):
'''

[rules.TRF046]
description = "forward must not write module attributes; modules are stateless in forward."
default_enabled = true
# No `cutoff_date`: 26 findings over 13 pre-existing models, mostly rotary/position caches written into
# `self` during `forward` (clvp, codegen, recurrent_gemma, seamless_m4t, udop, wav2vec2_conformer, ...).
# 9 are legacy `modeling_*.py`, 4 modular (qwen2_5_omni, qwen3_omni_moe, switch_transformers,
# wav2vec2_conformer). Making a module stateless changes runtime behavior, so each model needs its own PR
# with its tests re-run.
allowlist_models = ["clvp", "codegen", "ctrl", "fastspeech2_conformer", "nllb_moe", "qwen2_5_omni", "qwen3_omni_moe", "recurrent_gemma", "seamless_m4t", "switch_transformers", "trocr", "udop", "wav2vec2_conformer"]

[rules.TRF046.explanation]
what_it_does = "Flags assignments to self attributes in forward methods in modeling_*.py and modular_*.py."
why_bad = "State written during forward breaks batching, torch.compile, and reasoning about the module. Carried state is passed explicitly (cache objects, the generate loop); values that depend only on config or static shapes belong in __init__."
diff = '''
 def forward(self, hidden_states):
-    self.sequence_length = hidden_states.shape[1]
-    embeddings = self.compute_embeddings(self.sequence_length)
+    embeddings = self.compute_embeddings(hidden_states.shape[1])
'''

[rules.TRF047]
description = "Image/video processors are stateless: preprocess/_preprocess/post_process_* must not write self attributes."
default_enabled = true
allowlist_models = []

[rules.TRF047.explanation]
what_it_does = "Flags assignments to self attributes in preprocess, _preprocess, __call__ and post_process* methods in image_processing_*.py and video_processing_*.py."
why_bad = "A processor carrying state between calls breaks preprocess-many-then-postprocess batching: the second preprocess overwrites the state the first postprocess needs. Return the value or pass it through the method chain."
diff = '''
 def _preprocess(self, images, **kwargs):
-    self.original_sizes = [image.shape[-2:] for image in images]
+    original_sizes = [image.shape[-2:] for image in images]
     ...
+    return BatchFeature(data={"pixel_values": pixel_values, "original_sizes": original_sizes})
'''

[rules.TRF048]
description = "_tied_weights_keys must be a dict mapping target to source, not the pre-v5 list form."
default_enabled = true
allowlist_models = []

[rules.TRF048.explanation]
what_it_does = "Flags list/tuple/set literals in class-level _tied_weights_keys declarations."
why_bad = "v5 changed _tied_weights_keys to a dict mapping each tied target parameter to its source. The list form no longer says which parameter is the source, so tying, device_map computation and saving misbehave silently."
diff = '''
 class AcmeForCausalLM(AcmePreTrainedModel):
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
'''

[rules.TRF049]
description = "Weight initialization belongs in _init_weights, never in __init__ (meta-device init discards it)."
default_enabled = true
allowlist_models = ["convbert", "emu3", "esm", "granite_speech", "unispeech_sat"]

[rules.TRF049.explanation]
what_it_does = "Flags init calls in __init__ methods in modeling_*.py and modular_*.py: nn.init.* / init.* primitives and in-place initializers on own parameters (self.weight.data.normal_())."
why_bad = "Models instantiate on the meta device, so tensor values written in __init__ are discarded before loading: a parameter initialized only there holds random content when fine-tuning from scratch or after a meta-device reload. Allocate with torch.empty in __init__, initialize in _init_weights."
diff = '''
 class AcmeEmbeddings(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.position_embedding = nn.Parameter(torch.empty(config.num_positions, config.hidden_size))
-        nn.init.trunc_normal_(self.position_embedding, std=config.initializer_range)

 class AcmePreTrainedModel(PreTrainedModel):
     def _init_weights(self, module):
         super()._init_weights(module)
+        if isinstance(module, AcmeEmbeddings):
+            init.trunc_normal_(module.position_embedding, std=self.config.initializer_range)
'''

[rules.TRF050]
description = "One rotary module per model: attention classes must not instantiate their own rotary embedding."
default_enabled = true
allowlist_models = ["edgetam_video", "moshi", "recurrent_gemma", "sam2_video"]

[rules.TRF050.explanation]
what_it_does = "Flags calls to a *RotaryEmbedding class in the __init__ of classes whose name ends in Attention."
why_bad = "The Model owns a single rotary_emb, builds inv_freq once, and passes cos/sin down as position_embeddings. One rotary module per attention layer duplicates buffers, recomputes frequencies per layer, and diverges from the contract that attention receives position_embeddings."
diff = '''
 class AcmeAttention(nn.Module):
     def __init__(self, config, layer_idx):
         super().__init__()
-        self.rotary_emb = AcmeRotaryEmbedding(config)

 class AcmeModel(AcmePreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
+        self.rotary_emb = AcmeRotaryEmbedding(config)
'''

[rules.TRF051]
description = "Modeling code must not branch on _attn_implementation; the attention interface owns dispatch."
default_enabled = true
# No `cutoff_date`: 10 findings over 9 models, in two shapes that need a case-by-case call with the
# attention-interface owner. Legacy dispatch in plain `modeling_*.py`: decision_transformer, falcon, gpt2.
# Backend-conditional tensor handling in modular MLA/DSA-style models the interface cannot express yet:
# axk2, deepseek_v32, diffusion_gemma, glm_moe_dsa, minimax_m3_vl, qwen2_5_omni. Follow-ups either way.
allowlist_models = ["axk2", "decision_transformer", "deepseek_v32", "diffusion_gemma", "falcon", "glm_moe_dsa", "gpt2", "minimax_m3_vl", "qwen2_5_omni", "glm5_next", "hy_v4"]

[rules.TRF051.explanation]
what_it_does = "Flags comparisons against a _attn_implementation attribute in modeling_*.py and modular_*.py."
why_bad = "Backend dispatch belongs to ALL_ATTENTION_FUNCTIONS.get_interface, and backend-conditional tensor munging (padding, reshaping) belongs in the shared wrappers under integrations/. Inline branching keeps the model body kernel-aware and breaks when new backends register."
diff = '''
-if self.config._attn_implementation == "flash_attention_2":
-    attn_output = flash_path(query_states, key_states, value_states)
-else:
-    attn_output = eager_path(query_states, key_states, value_states)
+attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(self.config._attn_implementation, eager_attention_forward)
+attn_output, attn_weights = attention_interface(self, query_states, key_states, value_states, ...)
'''

[rules.TRF052]
description = "No *_ATTENTION_CLASSES dispatch dicts; attention backends route through the attention interface."
default_enabled = true
allowlist_models = ["bark", "data2vec", "falcon", "git", "gpt_neo", "gptj", "sam", "superglue"]

[rules.TRF052.explanation]
what_it_does = "Flags module-level assignments to names ending in _ATTENTION_CLASSES in modeling_*.py and modular_*.py."
why_bad = "Per-backend attention classes picked from a dict are the pre-interface idiom: near-identical classes drift apart, and hub attention kernels registered into ALL_ATTENTION_FUNCTIONS never reach them. One attention class dispatching through the interface replaces the dict; do not propagate it from a legacy parent."
diff = '''
-ACME_ATTENTION_CLASSES = {
-    "eager": AcmeAttention,
-    "flash_attention_2": AcmeFlashAttention2,
-    "sdpa": AcmeSdpaAttention,
-}
+class AcmeAttention(nn.Module):
+    def forward(self, ...):
+        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(self.config._attn_implementation, eager_attention_forward)
'''

[rules.TRF053]
description = "No manual label shifting in modeling code; self.loss_function owns it."
default_enabled = true
# No `cutoff_date`: 37 findings over 16 pre-existing models that slice `shift_logits` / `shift_labels` by
# hand, 12 of them in legacy `modeling_*.py` and 3 modular (csm, gemma3, modernbert_decoder). Each move to
# `self.loss_function` has to be checked against that model's own loss -- encoder-decoder models pass
# already-shifted labels -- so they go one model per PR.
allowlist_models = ["blip", "blip_2", "clvp", "csm", "gemma3", "git", "gpt2", "granite_speech", "imagegpt", "llama4", "mamba", "modernbert_decoder", "moshi", "openai", "qwen2_audio", "xlstm"]

[rules.TRF053.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags assignments that build shift_logits/shift_labels (and shifted_ variants) by slicing, as in labels[..., 1:]. Receiving already-shifted labels (shift_labels = kwargs.pop(\"shift_labels\", labels)) is the correct idiom and is not flagged."
why_bad = "self.loss_function shifts labels itself, so pre-shifting in modeling code trains on doubly-shifted targets or forces a bespoke loss path. Decoder-only models pass raw labels and let the loss shift them. Encoder-decoder models are the mirror case: their labels are already shifted by the prepended decoder start token, so they pass shift_labels=labels to stop the loss shifting again."
diff = '''
 if labels is not None:
-    shift_logits = logits[..., :-1, :].contiguous()
-    shift_labels = labels[..., 1:].contiguous()
-    loss = nn.functional.cross_entropy(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
+    # decoder-only: labels are unshifted, the loss shifts them
+    loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size)
+    # encoder-decoder: labels are already shifted, hand them over as shift_labels
+    loss = self.loss_function(logits=logits, labels=labels, shift_labels=labels, vocab_size=self.config.vocab_size)
'''

[rules.TRF054]
deprecated = true
description = "Removed in 0.1.4: flagged self.<media>_token_id assignments in processor __init__, too noisy on legitimate processor code."

[rules.TRF055]
description = "`config` on a PreTrainedModel subclass must be an annotation (`config: SomeConfig`), not an assignment (`config = SomeConfig`)."
default_enabled = true
allowlist_models = []

[rules.TRF055.explanation]
what_it_does = "Flags `config = SomeConfig` as a class attribute on PreTrainedModel subclasses in modeling_*.py and modular_*.py."
why_bad = "`PreTrainedModel.__init_subclass__` derives `config_class` from a `config` **annotation**, via `inspect.get_annotations(cls)`. An assignment creates a stray class attribute that call does not report, so the subclass silently keeps the parent's `config_class`. A bare annotation has no runtime value, creates no attribute, and is picked up correctly."
diff = '''
 class Gemma4VisionModel(Gemma4PreTrainedModel):
     """The Gemma 4 Vision Encoder."""
-    config = Gemma4VisionConfig
+    config: Gemma4VisionConfig
'''

[rules.TRF056]
description = "forward() must not materialize tensors with .item() or .tolist()."
default_enabled = true
# `kimi_k25` was dropped: no `.item()`/`.tolist()` left in its forward, verified against its source.
allowlist_models = ["hunyuan_vl"]
cutoff_date = "2026-06-20"

[rules.TRF056.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags `.item()` and `.tolist()` calls inside any `forward`. A `.tolist()` feeding the split-size argument of `split(...)` is exempt: torch.split needs Python ints."
why_bad = "Both calls read a tensor back to the host, which dynamo cannot trace, so the graph breaks."
diff = '''
-        for grid, item in zip(grid_thw.tolist(), split_items):
-            _, height, width = grid
-            merged.append(self.patch_merger(item, size=(height, width)))
+        for grid, item in zip(grid_thw, split_items):
+            merged.append(self.patch_merger(item, size=(grid[1], grid[2])))
'''

[rules.TRF057]
description = "Public model, config, output and processor classes and their public methods must carry @auto_docstring."
default_enabled = true
# no contribution date
# no contribution date
# no contribution date
allowlist_models = ["kosmos2", "timm_backbone", "timm_wrapper", "vitpose_backbone", "wav2vec2_with_lm", "x_clip"]
cutoff_date = "2026-06-20"

[rules.TRF057.explanation]
what_it_does = "Checks `@auto_docstring` on public `PreTrainedModel` subclasses (`<Model>PreTrainedModel`, `<Model>Model`, `<Model>For<Task>`, backbones), `PreTrainedConfig` subclasses, `ModelOutput` subclasses, image processors and `ProcessorMixin` subclasses, and on their public methods `forward`, `get_image_features`, `get_video_features`, `get_audio_features`, `get_text_features`, `preprocess` and `__call__`. A `modular_*.py` is checked against the files generated from it."
why_bad = "Without it a class ships with no intro and no parameter documentation, and a method with no argument documentation, no `Returns` section and no usage example -- all of which then have to be hand-written per model instead of coming from `auto_docstring.py`."
diff = '''
+@auto_docstring
 @dataclass
 class AcmeModelOutputWithPast(ModelOutput):
     logits: torch.FloatTensor | None = None

+@auto_docstring
 class AcmeForConditionalGeneration(AcmePreTrainedModel):
+    @auto_docstring
     def forward(self, input_ids, pixel_values=None, **kwargs):
         ...
'''

[rules.TRF058]
description = "Buffers must be declared as nn.Buffer attributes, not registered with register_buffer()."
default_enabled = true
allowlist_models = ["falcon_h1", "pp_doclayout_v2"]

[rules.TRF058.explanation]
what_it_does = "In modeling_*.py and modular_*.py, flags `register_buffer(\"<name>\", ...)` calls whose buffer name is a string literal, on any receiver (`self`, or another module such as `layer.mamba`). A computed name -- a variable or f-string, e.g. one buffer per layer inside a loop -- has no attribute-assignment equivalent and is exempt."
why_bad = "Since torch>=2.5 `nn.Buffer` registers a buffer through plain attribute assignment, like `nn.Parameter`. A buffer created by a method call only exists as a side effect of running `__init__`, so a modular file that wants to tweak one has to redefine the whole `__init__`. Assigned as an attribute, it can be inherited and overridden on its own."
diff = '''
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.register_buffer(
-            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
-        )
+        self.inv_freq = nn.Buffer(inv_freq, persistent=False)
+        self.position_ids = nn.Buffer(torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False)
'''

[rules.TRF059]
description = "Routed Experts modules tagged with `moe_tp_experts` must take hidden states, top-k indices and top-k weights in that order."
default_enabled = true
allowlist_models = ["llama4"]

[rules.TRF059.explanation]
what_it_does = "For model directories whose tensor-parallel plan assigns `moe_tp_experts`, checks that routed `*Experts` classes take hidden states, top-k indices and top-k routing weights as the first three positional `forward` arguments. Aliases such as `selected_experts` and `routing_weights` are accepted."
why_bad = "`MoeExpertsParallel` applies a gradient transform to positional argument 3. A different signature silently applies it to the wrong tensor or skips it."
diff = '''
 class AcmeExperts(nn.Module):
-    def forward(self, hidden_states):
+    def forward(self, hidden_states, top_k_index, top_k_weights):
         ...
'''