diff --git a/.cursor/rules/mod-000a-reusable-layers-belong-in-nn.mdc b/.cursor/rules/mod-000a-reusable-layers-belong-in-nn.mdc
index d1f5d544bb..3e9142cdd2 100644
--- a/.cursor/rules/mod-000a-reusable-layers-belong-in-nn.mdc
+++ b/.cursor/rules/mod-000a-reusable-layers-belong-in-nn.mdc
@@ -1,20 +1,19 @@
 ---
-description: Reusable layers and building blocks should be placed in physicsnemo/nn, not physicsnemo/models. Examples include FullyConnected, attention layers, and UNetBlock.
+description: Reusable layers and building blocks should be placed in physicsnemo/nn/module (exposed via physicsnemo.nn), not physicsnemo/models. Examples include FullyConnected, attention layers, and UNetBlock.
 alwaysApply: false
 ---
 
-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.
+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/module (and be re-exported through 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`,
+should go into `physicsnemo/nn/module`. 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:
+Implementations live in the `module/` subpackage but are re-exported from
+`physicsnemo/nn/__init__.py` so the public import path stays stable:
 
 ```python
 from physicsnemo.nn import MyLayer
@@ -33,13 +32,13 @@ and promotes code reuse across different models.
 **Example:**
 
 ```python
-# Good: Reusable layer in physicsnemo/nn/attention.py
+# Good: Reusable layer in physicsnemo/nn/module/attention_layers.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
+from physicsnemo.nn.module.attention_layers import MultiHeadAttention
 
 # Good: Example-specific layer in examples/weather/utils/nn.py
 class WeatherSpecificLayer(Module):
@@ -53,6 +52,6 @@ class WeatherSpecificLayer(Module):
 # WRONG: Reusable layer placed in physicsnemo/models/
 # File: physicsnemo/models/attention.py
 class MultiHeadAttention(Module):
-    """Should be in physicsnemo/nn/ not physicsnemo/models/"""
+    """Should be in physicsnemo/nn/module/ not physicsnemo/models/"""
     pass
 ```
diff --git a/.cursor/rules/mod-000b-complete-models-belong-in-models.mdc b/.cursor/rules/mod-000b-complete-models-belong-in-models.mdc
index 889bb4aae3..a817b57756 100644
--- a/.cursor/rules/mod-000b-complete-models-belong-in-models.mdc
+++ b/.cursor/rules/mod-000b-complete-models-belong-in-models.mdc
@@ -1,5 +1,5 @@
 ---
-description: Complete models composed of multiple layers should be placed in physicsnemo/models, not physicsnemo/nn. These are domain-specific or modality-specific models.
+description: Complete models composed of multiple layers should be placed in physicsnemo/models, not physicsnemo/nn/module (exposed via physicsnemo.nn). These are domain-specific or modality-specific models.
 alwaysApply: false
 ---
 
@@ -46,9 +46,9 @@ from physicsnemo.models.transformer import TransformerModel
 **Anti-pattern:**
 
 ```python
-# WRONG: Complete model placed in physicsnemo/nn/
-# File: physicsnemo/nn/transformer.py
+# WRONG: Complete model placed in physicsnemo/nn/module/
+# File: physicsnemo/nn/module/transformer.py
 class TransformerModel(Module):
-    """Should be in physicsnemo/models/ not physicsnemo/nn/"""
+    """Should be in physicsnemo/models/ not physicsnemo/nn/module/"""
     pass
 ```
diff --git a/.cursor/rules/mod-002a-experimental-models-belong-in-experimental.mdc b/.cursor/rules/mod-002a-experimental-models-belong-in-experimental.mdc
index 44076619e5..0e2bacc571 100644
--- a/.cursor/rules/mod-002a-experimental-models-belong-in-experimental.mdc
+++ b/.cursor/rules/mod-002a-experimental-models-belong-in-experimental.mdc
@@ -1,5 +1,5 @@
 ---
-description: New model classes should start in physicsnemo/experimental/nn or physicsnemo/experimental/models during development, where backward compatibility is not guaranteed.
+description: New model classes should start in physicsnemo/experimental/nn/module or physicsnemo/experimental/models during development, where backward compatibility is not guaranteed.
 alwaysApply: false
 ---
 
@@ -10,15 +10,16 @@ When creating new model or layer classes, rule MOD-002a must be followed. Explic
 **Description:**
 
 For the vast majority of models, new classes are created either in
-`physicsnemo/experimental/nn` for reusable layers, or in
+`physicsnemo/experimental/nn/module` for reusable layers, or in
 `physicsnemo/experimental/models` for more complete models. The `experimental`
 folder is used to store models that are still under development (beta or alpha
 releases) during this stage, backward compatibility is not guaranteed.
 
 One exception is when the developer is highly confident that the model is
 sufficiently mature and applicable to many domains or use cases. In this case
-the model class can be created in the `physicsnemo/nn` or `physicsnemo/models`
-folders directly, and backward compatibility is guaranteed.
+the model class can be created in the `physicsnemo/nn/module` (exposed through
+`physicsnemo.nn`) or `physicsnemo/models` folders directly, and backward
+compatibility is guaranteed.
 
 Another exception is when the model class is highly specific to a single
 example. In this case, it may be acceptable to place it in a module specific to
@@ -26,8 +27,9 @@ the example code, such as `examples/<example_name>/utils/nn.py`.
 
 After staying in experimental for a sufficient amount of time (typically at
 least 1 release cycle), the model class can be promoted to production. It is
-then moved to the `physicsnemo/nn` or `physicsnemo/models` folders, based on
-whether it's a reusable layer or complete model (see MOD-000a and MOD-000b).
+then moved to the `physicsnemo/nn/module` or `physicsnemo/models` folders,
+based on whether it's a reusable layer or complete model (see MOD-000a and
+MOD-000b).
 
 **Note:** Per MOD-008a, MOD-008b, and MOD-008c, it is forbidden to move a model
 out of the experimental stage/directory without the required CI tests.
diff --git a/.cursor/rules/mod-002b-add-deprecation-warnings-to-model.mdc b/.cursor/rules/mod-002b-add-deprecation-warnings-to-model.mdc
index 45e7f66883..d4b2303341 100644
--- a/.cursor/rules/mod-002b-add-deprecation-warnings-to-model.mdc
+++ b/.cursor/rules/mod-002b-add-deprecation-warnings-to-model.mdc
@@ -9,10 +9,11 @@ When deprecating a model class, rule MOD-002b must be followed. Explicitly refer
 
 **Description:**
 
-For a model class in the pre-deprecation stage in `physicsnemo/nn` or
-`physicsnemo/models`, the developer should start planning its deprecation. This
-is done by adding a warning message to the model class, indicating that the
-model class is deprecated and will be removed in a future release.
+For a model class in the pre-deprecation stage in `physicsnemo/nn/module`
+(exposed via `physicsnemo.nn`) or `physicsnemo/models`, the developer should
+start planning its deprecation. This is done by adding a warning message to the
+model class, indicating that the model class is deprecated and will be removed
+in a future release.
 
 The warning message should be a clear and concise message that explains why the
 model class is being deprecated and what the user should do instead. The
diff --git a/.cursor/rules/mod-007a-cannot-add-required-parameters-without-defaults.mdc b/.cursor/rules/mod-007a-cannot-add-required-parameters-without-defaults.mdc
index 21e0cd8789..36cf4292a3 100644
--- a/.cursor/rules/mod-007a-cannot-add-required-parameters-without-defaults.mdc
+++ b/.cursor/rules/mod-007a-cannot-add-required-parameters-without-defaults.mdc
@@ -9,9 +9,9 @@ When adding parameters to production models, rule MOD-007a must be strictly foll
 
 **Description:**
 
-For any model in `physicsnemo/nn` or `physicsnemo/models`, adding new required
-parameters (parameters without default values) to `__init__` or any public
-method is strictly forbidden. This breaks backward compatibility.
+For any model in `physicsnemo/nn/module` or `physicsnemo/models`, adding new
+required parameters (parameters without default values) to `__init__` or any
+public method is strictly forbidden. This breaks backward compatibility.
 
 New parameters must have default values to ensure existing code and checkpoints
 continue to work. If a new parameter is truly required, increment the model
diff --git a/.cursor/rules/mod-007b-cannot-remove-or-rename-parameters-without-compat-mapper.mdc b/.cursor/rules/mod-007b-cannot-remove-or-rename-parameters-without-compat-mapper.mdc
index 9862e0e6ee..405138f4fb 100644
--- a/.cursor/rules/mod-007b-cannot-remove-or-rename-parameters-without-compat-mapper.mdc
+++ b/.cursor/rules/mod-007b-cannot-remove-or-rename-parameters-without-compat-mapper.mdc
@@ -9,8 +9,9 @@ When removing or renaming parameters in production models, rule MOD-007b must be
 
 **Description:**
 
-For any model in `physicsnemo/nn` or `physicsnemo/models`, removing or renaming
-parameters is strictly forbidden without proper backward compatibility support.
+For any model in `physicsnemo/nn/module` or `physicsnemo/models`, removing or
+renaming parameters is strictly forbidden without proper backward compatibility
+support.
 
 If a parameter must be renamed or removed, the developer must:
 1. Increment `__model_checkpoint_version__`
diff --git a/.cursor/rules/mod-007c-cannot-change-return-types-of-public-methods.mdc b/.cursor/rules/mod-007c-cannot-change-return-types-of-public-methods.mdc
index dd347c0ab6..b898cacfb5 100644
--- a/.cursor/rules/mod-007c-cannot-change-return-types-of-public-methods.mdc
+++ b/.cursor/rules/mod-007c-cannot-change-return-types-of-public-methods.mdc
@@ -9,9 +9,9 @@ When modifying public method return types, rule MOD-007c must be strictly follow
 
 **Description:**
 
-For any model in `physicsnemo/nn` or `physicsnemo/models`, changing the return
-type of any public method (including `forward`) is strictly forbidden. This
-includes:
+For any model in `physicsnemo/nn/module` or `physicsnemo/models`, changing the
+return type of any public method (including `forward`) is strictly forbidden.
+This includes:
 - Changing from returning a single value to returning a tuple
 - Changing from a tuple to a single value
 - Changing the number of elements in a returned tuple
diff --git a/.cursor/rules/mod-008a-model-missing-constructor-attributes-tests.mdc b/.cursor/rules/mod-008a-model-missing-constructor-attributes-tests.mdc
index b188c133ae..f8512ceb6c 100644
--- a/.cursor/rules/mod-008a-model-missing-constructor-attributes-tests.mdc
+++ b/.cursor/rules/mod-008a-model-missing-constructor-attributes-tests.mdc
@@ -9,9 +9,9 @@ When creating tests for models, rule MOD-008a must be followed. Explicitly refer
 
 **Description:**
 
-Every model in `physicsnemo/nn` or `physicsnemo/models` must have tests that
-verify model instantiation and all public attributes (excluding buffers and
-parameters).
+Every model in `physicsnemo/nn/module` or `physicsnemo/models` must have tests
+that verify model instantiation and all public attributes (excluding buffers
+and parameters).
 
 These tests should:
 - Use `pytest` parameterization to test at least 2 configurations
diff --git a/CODING_STANDARDS/MODELS_IMPLEMENTATION.md b/CODING_STANDARDS/MODELS_IMPLEMENTATION.md
index 0000698f4e..0d40c54c9a 100644
--- a/CODING_STANDARDS/MODELS_IMPLEMENTATION.md
+++ b/CODING_STANDARDS/MODELS_IMPLEMENTATION.md
@@ -102,11 +102,10 @@ This document is structured in two main sections:
 **Description:**
 
 Reusable layers that are the building blocks of more complex architectures
-should go into `physicsnemo/nn`. Those include for instance `FullyConnected`,
+should go into `physicsnemo/nn/module`. 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:
+Implementations live in the `module/` subpackage but are re-exported from
+`physicsnemo/nn/__init__.py` so the public import path stays stable:
 
 ```python
 from physicsnemo.nn import MyLayer
@@ -125,13 +124,13 @@ and promotes code reuse across different models.
 **Example:**
 
 ```python
-# Good: Reusable layer in physicsnemo/nn/attention.py
+# Good: Reusable layer in physicsnemo/nn/module/attention_layers.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
+from physicsnemo.nn.module.attention_layers import MultiHeadAttention
 
 # Good: Example-specific layer in examples/weather/utils/nn.py
 class WeatherSpecificLayer(Module):
@@ -145,7 +144,7 @@ class WeatherSpecificLayer(Module):
 # WRONG: Reusable layer placed in physicsnemo/models/
 # File: physicsnemo/models/attention.py
 class MultiHeadAttention(Module):
-    """Should be in physicsnemo/nn/ not physicsnemo/models/"""
+    """Should be in physicsnemo/nn/module/ not physicsnemo/models/"""
     pass
 ```
 
@@ -192,10 +191,10 @@ from physicsnemo.models.transformer import TransformerModel
 **Anti-pattern:**
 
 ```python
-# WRONG: Complete model placed in physicsnemo/nn/
-# File: physicsnemo/nn/transformer.py
+# WRONG: Complete model placed in physicsnemo/nn/module/
+# File: physicsnemo/nn/module/transformer.py
 class TransformerModel(Module):
-    """Should be in physicsnemo/models/ not physicsnemo/nn/"""
+    """Should be in physicsnemo/models/ not physicsnemo/nn/module/"""
     pass
 ```
 
@@ -248,15 +247,16 @@ class MyModel(nn.Module):
 **Description:**
 
 For the vast majority of models, new classes are created either in
-`physicsnemo/experimental/nn` for reusable layers, or in
+`physicsnemo/experimental/nn/module` for reusable layers, or in
 `physicsnemo/experimental/models` for more complete models. The `experimental`
 folder is used to store models that are still under development (beta or alpha
 releases), where backward compatibility is not guaranteed.
 
 One exception is when the developer is highly confident that the model is
 sufficiently mature and applicable to many domains or use cases. In this case
-the model class can be created in the `physicsnemo/nn` or `physicsnemo/models`
-folders directly, and backward compatibility is guaranteed.
+the model class can be created in the `physicsnemo/nn/module` (exposed through
+`physicsnemo.nn`) or `physicsnemo/models` folders directly, and backward
+compatibility is guaranteed.
 
 Another exception is when the model class is highly specific to a single
 example. In this case, it may be acceptable to place it in a module specific to
@@ -264,9 +264,9 @@ the example code, such as `examples/<example_name>/utils/nn.py`.
 
 After staying in experimental for a sufficient amount of time (typically at
 least 1 release cycle), the model class can be promoted to production. It is
-then moved to the `physicsnemo/nn` or `physicsnemo/models` folders, based on
-whether it's a reusable layer (MOD-000a) or complete model (MOD-000b). During
-the production stage, backward compatibility is guaranteed.
+then moved to the `physicsnemo/nn/module` or `physicsnemo/models` folders,
+based on whether it's a reusable layer (MOD-000a) or complete model (MOD-000b).
+During the production stage, backward compatibility is guaranteed.
 
 **Note:** Per MOD-008a, MOD-008b, and MOD-008c, it is forbidden to move a model
 out of the experimental stage/directory without the required CI tests.
@@ -309,9 +309,10 @@ class BrandNewModel(Module):
 
 **Description:**
 
-For a model class being deprecated in `physicsnemo/nn` or `physicsnemo/models`,
-the developer must add warning messages indicating that the model class is
-deprecated and will be removed in a future release.
+For a model class being deprecated in `physicsnemo/nn/module` (exposed via
+`physicsnemo.nn`) or `physicsnemo/models`, the developer must add warning
+messages indicating that the model class is deprecated and will be removed in a
+future release.
 
 The warning message should be clear and concise, explaining why the model class
 is being deprecated and what the user should do instead. The deprecation message
@@ -1306,9 +1307,9 @@ def forward(
 
 **Description:**
 
-For any model in `physicsnemo/nn` or `physicsnemo/models`, adding new required
-parameters (parameters without default values) to `__init__` or any public
-method is strictly forbidden. This breaks backward compatibility.
+For any model in `physicsnemo/nn/module` or `physicsnemo/models`, adding new
+required parameters (parameters without default values) to `__init__` or any
+public method is strictly forbidden. This breaks backward compatibility.
 
 New parameters must have default values to ensure existing code and checkpoints
 continue to work. If a new parameter is truly required, increment the model
@@ -1363,8 +1364,9 @@ class MyModel(Module):
 
 **Description:**
 
-For any model in `physicsnemo/nn` or `physicsnemo/models`, removing or renaming
-parameters is strictly forbidden without proper backward compatibility support.
+For any model in `physicsnemo/nn/module` or `physicsnemo/models`, removing or
+renaming parameters is strictly forbidden without proper backward compatibility
+support.
 
 If a parameter must be renamed or removed, the developer must:
 1. Increment `__model_checkpoint_version__`
@@ -1447,9 +1449,9 @@ class MyModel(Module):
 
 **Description:**
 
-For any model in `physicsnemo/nn` or `physicsnemo/models`, changing the return
-type of any public method (including `forward`) is strictly forbidden. This
-includes:
+For any model in `physicsnemo/nn/module` or `physicsnemo/models`, changing the
+return type of any public method (including `forward`) is strictly forbidden.
+This includes:
 - Changing from returning a single value to returning a tuple
 - Changing from a tuple to a single value
 - Changing the number of elements in a returned tuple
@@ -1504,9 +1506,9 @@ class MyModel(Module):
 
 **Description:**
 
-Every model in `physicsnemo/nn` or `physicsnemo/models` must have tests that
-verify model instantiation and all public attributes (excluding buffers and
-parameters).
+Every model in `physicsnemo/nn/module` or `physicsnemo/models` must have tests
+that verify model instantiation and all public attributes (excluding buffers
+and parameters).
 
 These tests should:
 - Use `pytest` parameterization to test at least 2 configurations
diff --git a/physicsnemo/core/__init__.py b/physicsnemo/core/__init__.py
index 0264df0f51..baedbeebc0 100644
--- a/physicsnemo/core/__init__.py
+++ b/physicsnemo/core/__init__.py
@@ -17,6 +17,7 @@
 from .meta import ModelMetaData
 from .module import Module
 from .registry import ModelRegistry
+from .function import Function
 from .version_check import check_version_spec
 
-__all__ = ["ModelMetaData", "Module", "ModelRegistry"]
+__all__ = ["ModelMetaData", "Module", "ModelRegistry", "Function"]
diff --git a/physicsnemo/core/benchmark.py b/physicsnemo/core/benchmark.py
new file mode 100644
index 0000000000..1189426adf
--- /dev/null
+++ b/physicsnemo/core/benchmark.py
@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+from typing import Any, Callable, Iterable, Tuple
+
+import torch
+
+class BenchmarkMixin:
+    """Defines the hooks needed by the benchmark and validation helpers."""
+
+    @classmethod
+    def make_inputs(cls) -> Iterable[Tuple[str, dict[str, Any], Tuple[Any, ...]]]:
+        """Yield (label, init_kwargs, run_args) tuples for benchmarking."""
+
+        raise NotImplementedError(
+            f"{cls.__name__}.make_inputs must be implemented by subclasses"
+        )
+
+    @classmethod
+    def reference_impl(cls, *args, **kwargs):
+        """Compute a reference output for correctness checks."""
+
+        raise NotImplementedError(
+            f"{cls.__name__}.reference_impl must be implemented by subclasses"
+        )
+
+    @classmethod
+    def check(cls, *args, **kwargs) -> None:
+        """Validate outputs against the reference implementation."""
+
+        raise NotImplementedError(
+            f"{cls.__name__}.check must be implemented by subclasses"
+        )
+
+    @classmethod
+    def benchmark(
+        cls,
+        *,
+        repeats: int = 10,
+        warmup: int = 1,
+    ) -> None:
+        """Collect benchmark data for downstream consumers.
+
+        Parameters
+        ----------
+        repeats
+            Number of timed iterations recorded per benchmark case.
+        warmup
+            Runs discarded before timings are collected. Warmups help amortize
+            one-time setup costs.
+        """
+
+        raise NotImplementedError(
+            f"Benchmarking not supported yet"
+        )
\ No newline at end of file
diff --git a/physicsnemo/core/function.py b/physicsnemo/core/function.py
new file mode 100644
index 0000000000..84bf189106
--- /dev/null
+++ b/physicsnemo/core/function.py
@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import torch
+from torch.autograd import Function as TorchAutogradFunction
+
+from physicsnemo.core.benchmark import BenchmarkMixin
+
+class Function(TorchAutogradFunction, BenchmarkMixin):
+    """Base class for PhysicsNeMo custom autograd functions."""
+    # Placeholder for utilities to bring in warp, fuser, etc.
\ No newline at end of file
diff --git a/physicsnemo/core/module.py b/physicsnemo/core/module.py
index bd0c444f07..2484192f53 100644
--- a/physicsnemo/core/module.py
+++ b/physicsnemo/core/module.py
@@ -33,6 +33,7 @@
 
 import torch
 
+from physicsnemo.core.benchmark import BenchmarkMixin
 from physicsnemo.core.filesystem import _download_cached, _get_fs
 from physicsnemo.core.meta import ModelMetaData
 from physicsnemo.core.registry import ModelRegistry
@@ -67,7 +68,7 @@ def _load_state_dict_with_logging(
     return missing_keys, unexpected_keys
 
 
-class Module(torch.nn.Module):
+class Module(torch.nn.Module, BenchmarkMixin):
     """The base class for all network models in PhysicsNeMo.
 
     This should be used as a direct replacement for torch.nn.module and provides
diff --git a/physicsnemo/nn/__init__.py b/physicsnemo/nn/__init__.py
index ba5a46f4e5..1d1fcec859 100644
--- a/physicsnemo/nn/__init__.py
+++ b/physicsnemo/nn/__init__.py
@@ -14,7 +14,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .activations import (
+# Make physicsnemo.nn.Module an available import like torch.nn.Module
+from physicsnemo.core import Module
+
+from .module.activations import (
     CappedGELU,
     CappedLeakyReLU,
     Identity,
@@ -22,17 +25,17 @@
     Stan,
     get_activation,
 )
-from .ball_query import BQWarp
-from .conv_layers import ConvBlock, CubeEmbedding
-from .dgm_layers import DGMLayer
-from .fourier_layers import (
+from .module.ball_query import BQWarp
+from .module.conv_layers import ConvBlock, CubeEmbedding
+from .module.dgm_layers import DGMLayer
+from .module.fourier_layers import (
     FourierFilter,
     FourierLayer,
     FourierMLP,
     GaborFilter,
     fourier_encode,
 )
-from .fully_connected_layers import (
+from .module.fully_connected_layers import (
     Conv1dFCLayer,
     Conv2dFCLayer,
     Conv3dFCLayer,
@@ -40,26 +43,26 @@
     ConvNdKernel1Layer,
     FCLayer,
 )
-from .kan_layers import KolmogorovArnoldNetwork
-from .mlp_layers import Mlp
-from .resample_layers import (
+from .module.kan_layers import KolmogorovArnoldNetwork
+from .module.mlp_layers import Mlp
+from .module.resample_layers import (
     DownSample2D,
     DownSample3D,
     UpSample2D,
     UpSample3D,
 )
-from .siren_layers import SirenLayer, SirenLayerType
-from .spectral_layers import (
+from .module.siren_layers import SirenLayer, SirenLayerType
+from .module.spectral_layers import (
     SpectralConv1d,
     SpectralConv2d,
     SpectralConv3d,
     SpectralConv4d,
 )
-from .transformer_layers import (
+from .module.transformer_layers import (
     DecoderLayer,
     EncoderLayer,
     FuserLayer,
     SwinTransformer,
 )
-from .weight_fact import WeightFactLinear
-from .weight_norm import WeightNormLinear
+from .module.weight_fact import WeightFactLinear
+from .module.weight_norm import WeightNormLinear
\ No newline at end of file
diff --git a/physicsnemo/nn/functional/__init__.py b/physicsnemo/nn/functional/__init__.py
new file mode 100644
index 0000000000..f48d21e283
--- /dev/null
+++ b/physicsnemo/nn/functional/__init__.py
@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
\ No newline at end of file
diff --git a/physicsnemo/nn/module/__init__.py b/physicsnemo/nn/module/__init__.py
new file mode 100644
index 0000000000..f48d21e283
--- /dev/null
+++ b/physicsnemo/nn/module/__init__.py
@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
\ No newline at end of file
diff --git a/physicsnemo/nn/activations.py b/physicsnemo/nn/module/activations.py
similarity index 100%
rename from physicsnemo/nn/activations.py
rename to physicsnemo/nn/module/activations.py
diff --git a/physicsnemo/nn/attention_layers.py b/physicsnemo/nn/module/attention_layers.py
similarity index 100%
rename from physicsnemo/nn/attention_layers.py
rename to physicsnemo/nn/module/attention_layers.py
diff --git a/physicsnemo/nn/ball_query.py b/physicsnemo/nn/module/ball_query.py
similarity index 98%
rename from physicsnemo/nn/ball_query.py
rename to physicsnemo/nn/module/ball_query.py
index bb759998d5..f9701d595b 100644
--- a/physicsnemo/nn/ball_query.py
+++ b/physicsnemo/nn/module/ball_query.py
@@ -26,7 +26,7 @@
 import torch.nn as nn
 from einops import rearrange
 
-from physicsnemo.nn.neighbors import radius_search
+from physicsnemo.nn.module.neighbors import radius_search
 
 
 class BQWarp(nn.Module):
diff --git a/physicsnemo/nn/conv_layers.py b/physicsnemo/nn/module/conv_layers.py
similarity index 100%
rename from physicsnemo/nn/conv_layers.py
rename to physicsnemo/nn/module/conv_layers.py
diff --git a/physicsnemo/nn/dgm_layers.py b/physicsnemo/nn/module/dgm_layers.py
similarity index 100%
rename from physicsnemo/nn/dgm_layers.py
rename to physicsnemo/nn/module/dgm_layers.py
diff --git a/physicsnemo/nn/drop.py b/physicsnemo/nn/module/drop.py
similarity index 100%
rename from physicsnemo/nn/drop.py
rename to physicsnemo/nn/module/drop.py
diff --git a/physicsnemo/nn/fft.py b/physicsnemo/nn/module/fft.py
similarity index 100%
rename from physicsnemo/nn/fft.py
rename to physicsnemo/nn/module/fft.py
diff --git a/physicsnemo/nn/fourier_layers.py b/physicsnemo/nn/module/fourier_layers.py
similarity index 100%
rename from physicsnemo/nn/fourier_layers.py
rename to physicsnemo/nn/module/fourier_layers.py
diff --git a/physicsnemo/nn/fully_connected_layers.py b/physicsnemo/nn/module/fully_connected_layers.py
similarity index 100%
rename from physicsnemo/nn/fully_connected_layers.py
rename to physicsnemo/nn/module/fully_connected_layers.py
diff --git a/physicsnemo/nn/fused_silu.py b/physicsnemo/nn/module/fused_silu.py
similarity index 100%
rename from physicsnemo/nn/fused_silu.py
rename to physicsnemo/nn/module/fused_silu.py
diff --git a/physicsnemo/nn/gnn_layers/__init__.py b/physicsnemo/nn/module/gnn_layers/__init__.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/__init__.py
rename to physicsnemo/nn/module/gnn_layers/__init__.py
diff --git a/physicsnemo/nn/gnn_layers/bsms.py b/physicsnemo/nn/module/gnn_layers/bsms.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/bsms.py
rename to physicsnemo/nn/module/gnn_layers/bsms.py
diff --git a/physicsnemo/nn/gnn_layers/distributed_graph.py b/physicsnemo/nn/module/gnn_layers/distributed_graph.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/distributed_graph.py
rename to physicsnemo/nn/module/gnn_layers/distributed_graph.py
diff --git a/physicsnemo/nn/gnn_layers/embedder.py b/physicsnemo/nn/module/gnn_layers/embedder.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/embedder.py
rename to physicsnemo/nn/module/gnn_layers/embedder.py
diff --git a/physicsnemo/nn/gnn_layers/graph.py b/physicsnemo/nn/module/gnn_layers/graph.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/graph.py
rename to physicsnemo/nn/module/gnn_layers/graph.py
diff --git a/physicsnemo/nn/gnn_layers/graph_types.py b/physicsnemo/nn/module/gnn_layers/graph_types.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/graph_types.py
rename to physicsnemo/nn/module/gnn_layers/graph_types.py
diff --git a/physicsnemo/nn/gnn_layers/mesh_edge_block.py b/physicsnemo/nn/module/gnn_layers/mesh_edge_block.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/mesh_edge_block.py
rename to physicsnemo/nn/module/gnn_layers/mesh_edge_block.py
diff --git a/physicsnemo/nn/gnn_layers/mesh_graph_decoder.py b/physicsnemo/nn/module/gnn_layers/mesh_graph_decoder.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/mesh_graph_decoder.py
rename to physicsnemo/nn/module/gnn_layers/mesh_graph_decoder.py
diff --git a/physicsnemo/nn/gnn_layers/mesh_graph_encoder.py b/physicsnemo/nn/module/gnn_layers/mesh_graph_encoder.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/mesh_graph_encoder.py
rename to physicsnemo/nn/module/gnn_layers/mesh_graph_encoder.py
diff --git a/physicsnemo/nn/gnn_layers/mesh_graph_mlp.py b/physicsnemo/nn/module/gnn_layers/mesh_graph_mlp.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/mesh_graph_mlp.py
rename to physicsnemo/nn/module/gnn_layers/mesh_graph_mlp.py
diff --git a/physicsnemo/nn/gnn_layers/mesh_node_block.py b/physicsnemo/nn/module/gnn_layers/mesh_node_block.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/mesh_node_block.py
rename to physicsnemo/nn/module/gnn_layers/mesh_node_block.py
diff --git a/physicsnemo/nn/gnn_layers/utils.py b/physicsnemo/nn/module/gnn_layers/utils.py
similarity index 100%
rename from physicsnemo/nn/gnn_layers/utils.py
rename to physicsnemo/nn/module/gnn_layers/utils.py
diff --git a/physicsnemo/nn/interpolation.py b/physicsnemo/nn/module/interpolation.py
similarity index 100%
rename from physicsnemo/nn/interpolation.py
rename to physicsnemo/nn/module/interpolation.py
diff --git a/physicsnemo/nn/kan_layers.py b/physicsnemo/nn/module/kan_layers.py
similarity index 100%
rename from physicsnemo/nn/kan_layers.py
rename to physicsnemo/nn/module/kan_layers.py
diff --git a/physicsnemo/nn/layer_norm.py b/physicsnemo/nn/module/layer_norm.py
similarity index 100%
rename from physicsnemo/nn/layer_norm.py
rename to physicsnemo/nn/module/layer_norm.py
diff --git a/physicsnemo/nn/mlp_layers.py b/physicsnemo/nn/module/mlp_layers.py
similarity index 100%
rename from physicsnemo/nn/mlp_layers.py
rename to physicsnemo/nn/module/mlp_layers.py
diff --git a/physicsnemo/nn/neighbors/__init__.py b/physicsnemo/nn/module/neighbors/__init__.py
similarity index 100%
rename from physicsnemo/nn/neighbors/__init__.py
rename to physicsnemo/nn/module/neighbors/__init__.py
diff --git a/physicsnemo/nn/neighbors/_knn/__init__.py b/physicsnemo/nn/module/neighbors/_knn/__init__.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_knn/__init__.py
rename to physicsnemo/nn/module/neighbors/_knn/__init__.py
diff --git a/physicsnemo/nn/neighbors/_knn/_cuml_impl.py b/physicsnemo/nn/module/neighbors/_knn/_cuml_impl.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_knn/_cuml_impl.py
rename to physicsnemo/nn/module/neighbors/_knn/_cuml_impl.py
diff --git a/physicsnemo/nn/neighbors/_knn/_scipy_impl.py b/physicsnemo/nn/module/neighbors/_knn/_scipy_impl.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_knn/_scipy_impl.py
rename to physicsnemo/nn/module/neighbors/_knn/_scipy_impl.py
diff --git a/physicsnemo/nn/neighbors/_knn/_torch_impl.py b/physicsnemo/nn/module/neighbors/_knn/_torch_impl.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_knn/_torch_impl.py
rename to physicsnemo/nn/module/neighbors/_knn/_torch_impl.py
diff --git a/physicsnemo/nn/neighbors/_knn/knn.py b/physicsnemo/nn/module/neighbors/_knn/knn.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_knn/knn.py
rename to physicsnemo/nn/module/neighbors/_knn/knn.py
diff --git a/physicsnemo/nn/neighbors/_radius_search/__init__.py b/physicsnemo/nn/module/neighbors/_radius_search/__init__.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_radius_search/__init__.py
rename to physicsnemo/nn/module/neighbors/_radius_search/__init__.py
diff --git a/physicsnemo/nn/neighbors/_radius_search/_torch_impl.py b/physicsnemo/nn/module/neighbors/_radius_search/_torch_impl.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_radius_search/_torch_impl.py
rename to physicsnemo/nn/module/neighbors/_radius_search/_torch_impl.py
diff --git a/physicsnemo/nn/neighbors/_radius_search/_warp_impl.py b/physicsnemo/nn/module/neighbors/_radius_search/_warp_impl.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_radius_search/_warp_impl.py
rename to physicsnemo/nn/module/neighbors/_radius_search/_warp_impl.py
diff --git a/physicsnemo/nn/neighbors/_radius_search/kernels.py b/physicsnemo/nn/module/neighbors/_radius_search/kernels.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_radius_search/kernels.py
rename to physicsnemo/nn/module/neighbors/_radius_search/kernels.py
diff --git a/physicsnemo/nn/neighbors/_radius_search/radius_search.py b/physicsnemo/nn/module/neighbors/_radius_search/radius_search.py
similarity index 100%
rename from physicsnemo/nn/neighbors/_radius_search/radius_search.py
rename to physicsnemo/nn/module/neighbors/_radius_search/radius_search.py
diff --git a/physicsnemo/nn/resample_layers.py b/physicsnemo/nn/module/resample_layers.py
similarity index 100%
rename from physicsnemo/nn/resample_layers.py
rename to physicsnemo/nn/module/resample_layers.py
diff --git a/physicsnemo/nn/sdf.py b/physicsnemo/nn/module/sdf.py
similarity index 100%
rename from physicsnemo/nn/sdf.py
rename to physicsnemo/nn/module/sdf.py
diff --git a/physicsnemo/nn/siren_layers.py b/physicsnemo/nn/module/siren_layers.py
similarity index 100%
rename from physicsnemo/nn/siren_layers.py
rename to physicsnemo/nn/module/siren_layers.py
diff --git a/physicsnemo/nn/spectral_layers.py b/physicsnemo/nn/module/spectral_layers.py
similarity index 100%
rename from physicsnemo/nn/spectral_layers.py
rename to physicsnemo/nn/module/spectral_layers.py
diff --git a/physicsnemo/nn/transformer_decoder.py b/physicsnemo/nn/module/transformer_decoder.py
similarity index 100%
rename from physicsnemo/nn/transformer_decoder.py
rename to physicsnemo/nn/module/transformer_decoder.py
diff --git a/physicsnemo/nn/transformer_layers.py b/physicsnemo/nn/module/transformer_layers.py
similarity index 100%
rename from physicsnemo/nn/transformer_layers.py
rename to physicsnemo/nn/module/transformer_layers.py
diff --git a/physicsnemo/nn/utils/__init__.py b/physicsnemo/nn/module/utils/__init__.py
similarity index 100%
rename from physicsnemo/nn/utils/__init__.py
rename to physicsnemo/nn/module/utils/__init__.py
diff --git a/physicsnemo/nn/utils/patch_embed.py b/physicsnemo/nn/module/utils/patch_embed.py
similarity index 100%
rename from physicsnemo/nn/utils/patch_embed.py
rename to physicsnemo/nn/module/utils/patch_embed.py
diff --git a/physicsnemo/nn/utils/shift_window_mask.py b/physicsnemo/nn/module/utils/shift_window_mask.py
similarity index 100%
rename from physicsnemo/nn/utils/shift_window_mask.py
rename to physicsnemo/nn/module/utils/shift_window_mask.py
diff --git a/physicsnemo/nn/utils/utils.py b/physicsnemo/nn/module/utils/utils.py
similarity index 100%
rename from physicsnemo/nn/utils/utils.py
rename to physicsnemo/nn/module/utils/utils.py
diff --git a/physicsnemo/nn/utils/weight_init.py b/physicsnemo/nn/module/utils/weight_init.py
similarity index 100%
rename from physicsnemo/nn/utils/weight_init.py
rename to physicsnemo/nn/module/utils/weight_init.py
diff --git a/physicsnemo/nn/weight_fact.py b/physicsnemo/nn/module/weight_fact.py
similarity index 100%
rename from physicsnemo/nn/weight_fact.py
rename to physicsnemo/nn/module/weight_fact.py
diff --git a/physicsnemo/nn/weight_norm.py b/physicsnemo/nn/module/weight_norm.py
similarity index 100%
rename from physicsnemo/nn/weight_norm.py
rename to physicsnemo/nn/module/weight_norm.py