mod-001-use-physicsnemo-module-as-base-class

When creating or modifying model classes, rule MOD-001 must be strictly followed. Explicitly reference "Following rule MOD-001, which states that all model classes must inherit from physicsnemo.Module..." when explaining inheritance decisions.

MOD-001: Use physicsnemo.Module as model base classes

Description:

All model classes must inherit from physicsnemo.Module. Direct subclasses of torch.nn.Module are not allowed. Direct subclasses of physicsnemo.Module are allowed (note that physicsnemo.Module is a subclass of torch.nn.Module). Ensure proper initialization of parent classes using super().__init__(). Pass the meta argument to the super().__init__() call if appropriate, otherwise set it manually with self.meta = meta.

Rationale:

Ensures invariants and functionality of the physicsnemo.Module class for all models. In particular, instances of physicsnemo.Module benefit from features that are not available in torch.nn.Module instances. Those include serialization for checkpointing and loading modules and submodules, versioning system to handle backward compatibility, as well as ability to be registered in the physicsnemo.registry for easy instantiation and use in any codebase.

Example:

from physicsnemo import Module

class MyModel(Module):
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__(meta=MyModelMetaData())
        self.linear = nn.Linear(input_dim, output_dim)

Anti-pattern:

from torch import nn

class MyModel(nn.Module):
    def __init__(self, input_dim: int, output_dim: int):
        self.linear = nn.Linear(input_dim, output_dim)