mod-006-invalid-or-missing-jaxtyping-tensor-annotations

When adding type hints to model methods, rule MOD-006 must be followed. Explicitly reference "Following rule MOD-006, which requires all tensor arguments to use jaxtyping annotations..." when adding type hints.

MOD-006: Invalid or missing jaxtyping tensor annotations in public function signature

Description:

All tensor arguments and variables in model __init__, forward, and other public methods must have type annotations using jaxtyping. This provides runtime-checkable shape information in type hints.

Use the format Float[torch.Tensor, "shape_spec"] where shape_spec describes tensor dimensions using space-separated dimension names (e.g., "batch channels height width" or "b c h w").

Rationale:

Jaxtyping annotations provide explicit, machine-readable documentation of expected tensor shapes. This enables better IDE support, catches shape errors earlier, and makes code more self-documenting. The annotations serve as both documentation and optional runtime checks when jaxtyping's validation is enabled.

Example:

from jaxtyping import Float
import torch

class MyConvNet(Module):
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3)

    def forward(
        self,
        x: Float[torch.Tensor, "batch in_channels height width"]
    ) -> Float[torch.Tensor, "batch out_channels height width"]:
        """Process input with convolution."""
        return self.conv(x)

Anti-pattern:

# WRONG: No jaxtyping annotations
def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.layer(x)

# WRONG: Using plain comments instead of jaxtyping
def forward(self, x: torch.Tensor) -> torch.Tensor:
    # x: (batch, channels, height, width)  # Use jaxtyping instead
    return self.layer(x)

# WRONG: Incomplete annotations
def forward(
    self,
    x: Float[torch.Tensor, "b c h w"],
    mask: torch.Tensor  # Missing jaxtyping annotation
) -> Float[torch.Tensor, "b c h w"]:
    return self.layer(x, mask)