TRF021

Scalar tensors must be filled on-device with torch.full((), …) instead of torch.tensor(…, device=…).

Default	Enabled
Scope	All models
Source	mlinter/trf021.py
Show in terminal	mlinter --rule TRF021

What it does

In modeling_*.py and modular_*.py, checks calls to torch.tensor(<value>, …, device=<non-cpu>) whose <value> provably evaluates to a Python scalar. Scalar-ness is resolved statically: numeric literals, arithmetic over them, torch.finfo/torch.iinfo fields, scalar-returning builtins and math.* calls, locals bound exactly once, self.<attr> assigned in the class body, and self.config.<field>/config.<field> whose annotation in the companion configuration file is int/float/bool (following attribute_map aliases). Fields that may also be sequences, such as eos_token_id: int | list[int] | None, and any expression that cannot be resolved are left alone. Construction-time methods (__init__, _init_weights, __post_init__, post_init) are exempt because they never run inside a capture region.

Why is this bad?

torch.tensor(<python scalar>, device=<accelerator>) materialises the value on the host and then issues a host-to-device copy. CUDA graph capture forbids that copy, so the model cannot be captured. torch.full((), <value>, dtype=…, device=…) fills the same 0-d tensor directly on-device with a capturable kernel and no synchronisation.

Example

 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)

Suppressing this rule

Add a # trf-ignore: TRF021 comment on the flagged line or the line directly above it. See Suppressing rules for whole-file directives and when a suppression is the wrong answer.