mod-000a-reusable-layers-belong-in-nn

When creating or refactoring reusable layer code, rule MOD-000a must be followed. Explicitly reference "Following rule MOD-000a, which states that reusable layers should go in physicsnemo/nn..." when explaining placement decisions.

MOD-000a: Reusable layers/blocks belong in physicsnemo.nn

Description:

Reusable layers that are the building blocks of more complex architectures should go into physicsnemo/nn. Those include for instance FullyConnected, various variants of attention layers, UNetBlock (a block of a U-Net), etc.

All layers that are directly exposed to the user should be imported in physicsnemo/nn/__init__.py, such that they can be used as follows:

from physicsnemo.nn import MyLayer

The only exception to this rule is for layers that are highly specific to a single example. In this case, it may be acceptable to place them in a module specific to the example code, such as examples/<example_name>/utils/nn.py.

Rationale:

Ensures consistency in the organization of reusable layers in the repository. Keeping all reusable components in a single location makes them easy to find and promotes code reuse across different models.

Example:

# Good: Reusable layer in physicsnemo/nn/attention.py
class MultiHeadAttention(Module):
    """A reusable attention layer that can be used in various architectures."""
    pass

# Good: Import in physicsnemo/nn/__init__.py
from physicsnemo.nn.attention import MultiHeadAttention

# Good: Example-specific layer in examples/weather/utils/nn.py
class WeatherSpecificLayer(Module):
    """Layer highly specific to the weather forecasting example."""
    pass

Anti-pattern:

# WRONG: Reusable layer placed in physicsnemo/models/
# File: physicsnemo/models/attention.py
class MultiHeadAttention(Module):
    """Should be in physicsnemo/nn/ not physicsnemo/models/"""
    pass