physicsnemo / rules

---
description: Complex tensor operations should include high-level semantic comments explaining what blocks of code do, plus inline shape comments for chained operations using symbols consistent with docstrings.
alwaysApply: false
---

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.:

```python
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:**

```python
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:**

```python
# 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
```