mod-005-invalid-or-missing-tensor-shape-validation

When implementing forward or public methods, rule MOD-005 must be followed. Explicitly reference "Following rule MOD-005, which requires tensor shape validation at the beginning of methods with torch.compiler.is_compiling() guard..." when adding validation code.

MOD-005: Invalid or missing tensor shape validation logic

Description:

All forward methods and other public methods that accept tensor arguments must validate tensor shapes at the beginning of the method. This rule applies to:

• Individual tensor arguments
• Containers of tensors (lists, tuples, dictionaries)

For containers, validate their length, required keys, and the shapes of contained tensors. Validation statements should be concise (ideally one check per argument). Error messages must follow the standardized format: "Expected tensor of shape (B, D) but got tensor of shape {actual_shape}".

To avoid interactions with torch.compile, all validation must be wrapped in a conditional check using torch.compiler.is_compiling(). Follow the "fail-fast" approach by validating inputs before any computation.

Rationale:

Early shape validation catches errors at the API boundary with clear, actionable error messages, making debugging significantly easier. Without validation, shape mismatches result in cryptic errors deep in the computation graph. The torch.compile guard ensures that validation overhead is eliminated in production compiled code while preserving debug-time safety.

Example:

def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Forward pass with shape validation."""
    ### Input validation
    # Skip validation when running under torch.compile for performance
    if not torch.compiler.is_compiling():
        # Extract expected dimensions
        B, C, H, W = x.shape if x.ndim == 4 else (None, None, None, None)

        # Validate x shape
        if x.ndim != 4:
            raise ValueError(
                f"Expected 4D input tensor (B, C, H, W), got {x.ndim}D tensor with shape {tuple(x.shape)}"
            )

        if C != self.in_channels:
            raise ValueError(
                f"Expected {self.in_channels} input channels, got {C} channels"
            )

        # Validate optional mask
        if mask is not None:
            if mask.shape != (B, H, W):
                raise ValueError(
                    f"Expected mask shape ({B}, {H}, {W}), got {tuple(mask.shape)}"
                )

    # Actual computation happens after validation
    return self._process(x, mask)

Anti-pattern:

# WRONG: No validation at all
def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.layer(x)  # Will fail with cryptic error if shape is wrong

# WRONG: Validation not guarded by torch.compiler.is_compiling()
def forward(self, x: torch.Tensor) -> torch.Tensor:
    if x.ndim != 4:  # Breaks torch.compile
        raise ValueError(f"Expected 4D tensor, got {x.ndim}D")
    return self.layer(x)

# WRONG: Validation after computation has started
def forward(self, x: torch.Tensor) -> torch.Tensor:
    h = self.layer1(x)  # Computation started
    if x.shape[1] != self.in_channels:  # Too late!
        raise ValueError(f"Wrong number of channels")
    return self.layer2(h)