mod-003k-add-high-level-comments-for-complex-tensor-operations

When writing model code with complex tensor operations, rule MOD-003k should be followed. Explicitly reference "Following rule MOD-003k, which recommends high-level comments for complex tensor operations..." when adding explanatory comments.

MOD-003k: Add high-level comments for complex tensor operations

Description:

Model code that involves complex tensor operations should include high-level comments that explain what blocks of code accomplish semantically. One-line comments every few lines of tensor operations is sufficient.

Comments should focus on high-level semantic explanations rather than low-level syntactic details. For example, use "Compute the encodings" instead of "Doing a concatenation followed by a linear projection, followed by a nonlinear activation". The goal is to give a high-level overview of what a block of tensor operations accomplishes.

When multiple tensor operations are chained, it is welcomed to add short inline comments with the tensor shapes of computed tensors, e.g.:

x = torch.cat([y, z], dim=1)  # (B, 2*C_in, H, W)

The symbols chosen in the comments should be consistent with the docstring (possibly shortened versions of dimension names for explicitness).

Rationale:

High-level comments make complex tensor manipulation code more understandable without cluttering it with excessive detail. Shape annotations help developers track tensor dimensions through complex operations, catching shape mismatches early. Consistency with docstring notation creates a unified mental model.

Example:

def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
    """Process input with context conditioning."""
    # Encode input features
    h = self.encoder(x)  # (B, C_enc, H, W)

    # Combine with context information
    c = self.context_proj(context)  # (B, C_enc)
    c = c[:, :, None, None].expand(-1, -1, h.shape[2], h.shape[3])  # (B, C_enc, H, W)
    h = torch.cat([h, c], dim=1)  # (B, 2*C_enc, H, W)

    # Apply attention mechanism
    h = self.attention(h)  # (B, 2*C_enc, H, W)

    # Decode to output
    out = self.decoder(h)  # (B, C_out, H, W)

    return out

Anti-pattern:

# WRONG: No comments for complex operations
def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
    h = self.encoder(x)
    c = self.context_proj(context)
    c = c[:, :, None, None].expand(-1, -1, h.shape[2], h.shape[3])
    h = torch.cat([h, c], dim=1)
    h = self.attention(h)
    out = self.decoder(h)
    return out

# WRONG: Too low-level, syntactic comments
def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
    # Pass x through encoder layer
    h = self.encoder(x)
    # Project context using linear layer
    c = self.context_proj(context)
    # Add two None dimensions and expand
    c = c[:, :, None, None].expand(-1, -1, h.shape[2], h.shape[3])
    # Concatenate h and c along dimension 1
    h = torch.cat([h, c], dim=1)
    # Apply attention
    h = self.attention(h)
    # Pass through decoder
    out = self.decoder(h)
    return out