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, 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 is this 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.

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.