Added BatchNorm2d, LayerNorm, InstanceNorm2d nodes.

__init__.py

@@ -146,9 +146,12 @@ from .modules.pytorch_wrapper.ptf_softmax import PtfSoftmax
 from .modules.pytorch_wrapper.ptf_softplus import PtfSoftplus
 from .modules.pytorch_wrapper.ptf_tanh import PtfTanh
 from .modules.pytorch_wrapper.ptn_avg_pool2d import PtnAvgPool2d
+from .modules.pytorch_wrapper.ptn_batch_norm2d import PtnBatchNorm2d
 from .modules.pytorch_wrapper.ptn_chained_model import PtnChainedModel
 from .modules.pytorch_wrapper.ptn_conv2d import PtnConv2d
 from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel
+from .modules.pytorch_wrapper.ptn_instance_norm2d import PtnInstanceNorm2d
+from .modules.pytorch_wrapper.ptn_layer_norm import PtnLayerNorm
 from .modules.pytorch_wrapper.ptn_linear import PtnLinear
 from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel
 from .modules.pytorch_wrapper.ptn_max_pool2d import PtnMaxPool2d
@@ -324,9 +327,12 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtfSoftplus": PtfSoftplus,
     "PtfTanh": PtfTanh,
     "PtnAvgPool2d": PtnAvgPool2d,
+    "PtnBatchNorm2d": PtnBatchNorm2d,
     "PtnChainedModel": PtnChainedModel,
     "PtnConv2d": PtnConv2d,
     "PtnConvModel": PtnConvModel,
+    "PtnInstanceNorm2d": PtnInstanceNorm2d,
+    "PtnLayerNorm": PtnLayerNorm,
     "PtnLinear": PtnLinear,
     "PtnLinearModel": PtnLinearModel,
     "PtnMaxPool2d": PtnMaxPool2d,
@@ -501,13 +507,16 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtfSoftmax": "Ptf Softmax",
     "PtfSoftplus": "Ptf Softplus",
     "PtfTanh": "Ptf Tanh",
-    "PtnAvgPool2d": "Ptn Avg Pool2d",
+    "PtnAvgPool2d": "Ptn Avg Pool 2d",
+    "PtnBatchNorm2d": "Ptn Batch Norm 2d",
     "PtnChainedModel": "Ptn Chained Model",
-    "PtnConv2d": "Ptn Conv2d",
+    "PtnConv2d": "Ptn Conv 2d",
     "PtnConvModel": "Ptn Conv Model",
+    "PtnInstanceNorm2d": "Ptn Instance Norm 2d",
+    "PtnLayerNorm": "Ptn Layer Norm",
     "PtnLinear": "Ptn Linear",
     "PtnLinearModel": "Ptn Linear Model",
-    "PtnMaxPool2d": "Ptn Max Pool2d",
+    "PtnMaxPool2d": "Ptn Max Pool 2d",
     "PtnPreAddChannelAxis": "Ptn Pre Add Channel Axis",
     "PtnPreFlatten": "Ptn Pre Flatten",
     "PtnResnetModel": "Ptn Resnet Model",

dev/scripts/init_py_generate.py

@@ -90,7 +90,8 @@ def format_display_name(class_name):
         ("MPL Scatter", "MPL Scatter Plot"),
         ("SNS Bar", "SNS Bar Chart"),
         ("SNS Line", "SNS Line Plot"),
-        ("SNS Scatter", "SNS Scatter Plot")
+        ("SNS Scatter", "SNS Scatter Plot"),
+        ("2d", " 2d"),
     ]
     for p in replace_pairs:
         final_result = final_result.replace(p[0], p[1])

docs/reference/node_reference.md

@@ -138,13 +138,16 @@
 ## PyTorch wrapper - Model
 | Node | Description |
 | --- | --- |
-| [Ptn Avg Pool](ptn_avg_pool2d.md) | An avgpool layer.   |
+| [Ptn Avg Pool 2d](ptn_avg_pool2d.md) | An avgpool layer.   |
+| [Ptn Batch Norm 2d](ptn_batch_norm2d.md) | A normalization model to normalize over the batch and spatial axes for each channel. |
 | [Ptn Chained Model](ptn_chained_model.md) | Constructs a chained PyTorch model. |
-| [Ptn Conv](ptn_conv2d.md) |     A convolutional model consisting of a single conv2d layer.   |
+| [Ptn Conv 2d](ptn_conv2d.md) |     A convolutional model consisting of a single conv2d layer.   |
 | [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers.   |
+| [Ptn Instance Norm 2d](ptn_instance_norm2d.md) | A normalization model to normalize elements over spatial axes within each channel for each sample. |
+| [Ptn Layer Norm](ptn_layer_norm.md) | A normalization model to normalize elements over specified axes. |
 | [Ptn Linear](ptn_linear.md) | A linear model consisting of a single dense layer.   |
 | [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers.   |
-| [Ptn Max Pool](ptn_max_pool2d.md) | A maxpool layer.   |
+| [Ptn Max Pool 2d](ptn_max_pool2d.md) | A maxpool layer.   |
 | [Ptn Pre Add Channel Axis](ptn_pre_add_channel_axis.md) | Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w) |
 | [Ptn Pre Flatten](ptn_pre_flatten.md) | Flattens the input tensor before processing the tensor in the specified model. |
 | [Ptn Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers.   |

docs/reference/ptn_avg_pool2d.md

@@ -1,4 +1,4 @@
-# Ptn Avg Pool
+# Ptn Avg Pool 2d
 An avgpool layer.  
 
     Args:

docs/reference/ptn_batch_norm2d.md

@@ -0,0 +1,30 @@
+# Ptn Batch Norm 2d
+A normalization model to normalize over the batch and spatial axes for each channel.
+
+Args:
+    num_features (int): Specifies the size of channel axis.
+        - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+    affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+    track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current batch statistics are always used for mean and std irrespective of this flag.
+        
+    momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher values give more weight to the current batch statistics, while lower values retain more past statistics.
+
+## Input
+| Name | Data type |
+|---|---|
+| num_features | Int |
+| affine | Boolean |
+| track_running_stats | Boolean |
+| momentum | Float |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptn_conv2d.md

@@ -1,4 +1,4 @@
-# Ptn Conv
+# Ptn Conv 2d
 A convolutional model consisting of a single conv2d layer.  
 
         Args:

docs/reference/ptn_instance_norm2d.md

@@ -0,0 +1,30 @@
+# Ptn Instance Norm 2d
+A normalization model to normalize elements over spatial axes within each channel for each sample.
+
+    Args:
+        num_features (int): Specifies the size of channel axis.
+            - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+        affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+        track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current sample data is always used for mean and std irrespective of this flag.
+         
+        momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher value gives more weight to the current statistics (mean & std).
+
+## Input
+| Name | Data type |
+|---|---|
+| num_features | Int |
+| affine | Boolean |
+| track_running_stats | Boolean |
+| momentum | Float |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptn_layer_norm.md

@@ -0,0 +1,32 @@
+# Ptn Layer Norm
+A normalization model to normalize elements over specified axes.
+
+    Args:
+        normalized_shape (list): Specifies the shape of elements to normalize.  
+            - For an input of shape `[8, 1024, 768]` (where axes represent `[batch_size, sequence_length, hidden_dim]`):
+            - `normalize_shape=[768]` normalizes within each token (across the hidden dimension).
+            - `normalize_shape=[1024, 768]` normalizes within each sample (across both sequence and token dimensions).  
+            - For an image input of shape `[8, 4, 256, 256]` (where axes represent `[batch_size, channel, height, width]`):
+            - `normalize_shape=[256, 256]` normalizes within each channel (across spatial dimensions).
+            - `normalize_shape=[4, 256, 256]` normalizes within each sample (across channels and spatial dimensions).
+
+        elementwise_affine (bool): If `True`, applies a learnable scaling factor to the normalized elements.
+
+        bias (bool): If `True` and elementwise_affine is also `True`, applies a learnable bias to the normalized elements.
+
+## Input
+| Name | Data type |
+|---|---|
+| normalized_shape | String |
+| elementwise_affine | Boolean |
+| bias | Boolean |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptn_max_pool2d.md

@@ -1,4 +1,4 @@
-# Ptn Max Pool
+# Ptn Max Pool 2d
 A maxpool layer.  
 
     Args:

modules/pytorch_wrapper/ptn_avg_pool2d.py

@@ -5,6 +5,7 @@ from .utils import str_to_dim
 
 class PtnAvgPool2d:
     """
+    Ptn Avg Pool 2d:
     An avgpool layer.  
 
         Args:

modules/pytorch_wrapper/ptn_batch_norm2d.py

@@ -0,0 +1,80 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+from .utils import str_to_dim
+
+class PtnBatchNorm2d:
+    """
+    Ptn Batch Norm 2d:
+    A normalization model to normalize over the batch and spatial axes for each channel.
+
+    Args:
+        num_features (int): Specifies the size of channel axis.
+            - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+        affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+        track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current batch statistics are always used for mean and std irrespective of this flag.
+            
+        momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher values give more weight to the current batch statistics, while lower values retain more past statistics.
+             
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "num_features": ("INT", {"default":1, "min":1, "max":1e6}),
+                "affine": ("BOOLEAN", {"default": True}),
+                "track_running_stats": ("BOOLEAN", {"default": True}),
+                "momentum": ("FLOAT", {"default":0.1, "min":0, "max":1, "step":1e-6}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          num_features: int,
+          affine: bool,
+          track_running_stats: bool,
+          momentum: float
+          ) -> Tuple[nn.Module]:
+        """
+        Instantiates a normalization model to normalize over the batch axis and spacial axes.
+
+        Args:
+            num_features (int): Specifies the size of channel axis.
+                - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+            affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+            track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current batch statistics are always used for mean and std irrespective of this flag.
+             
+            momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher values give more weight to the current batch statistics, while lower values retain more past statistics.
+            
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+
+        with torch.inference_mode(False):
+            model = nn.BatchNorm2d(
+                num_features=num_features,
+                affine=affine,
+                track_running_stats=track_running_stats,
+                momentum=momentum)
+                    
+        return (model,)

modules/pytorch_wrapper/ptn_conv2d.py

@@ -6,6 +6,7 @@ from .utils import str_to_dim
 
 class PtnConv2d:
     """
+    Ptn Conv 2d:
     A convolutional model consisting of a single conv2d layer.  
 
         Args:

modules/pytorch_wrapper/ptn_instance_norm2d.py

@@ -0,0 +1,80 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+from .utils import str_to_dim
+
+class PtnInstanceNorm2d:
+    """
+    Ptn Instance Norm 2d:
+    A normalization model to normalize elements over spatial axes within each channel for each sample.
+
+        Args:
+            num_features (int): Specifies the size of channel axis.
+                - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+            affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+            track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current sample data is always used for mean and std irrespective of this flag.
+             
+            momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher value gives more weight to the current statistics (mean & std).
+             
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "num_features": ("INT", {"default":1, "min":1, "max":1e6}),
+                "affine": ("BOOLEAN", {"default": True}),
+                "track_running_stats": ("BOOLEAN", {"default": True}),
+                "momentum": ("FLOAT", {"default":0.1, "min":0, "max":1, "step":1e-6}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          num_features: int,
+          affine: bool,
+          track_running_stats: bool,
+          momentum: float
+          ) -> Tuple[nn.Module]:
+        """
+        A normalization model to normalize elements over spatial axes within each channel for each sample.
+
+        Args:
+            num_features (int): Specifies the size of channel axis.
+                - For an image input of shape `[8, 4, 256, 256]`, specify 4.
+
+            affine (bool): If `True`, applies a learnable scaling factor and bias to the normalized elements.
+
+            track_running_stats (bool): If `True`, use exponential moving average (EMA) to store mean and std to be used in eval for normalization. For training, current sample data is always used for mean and std irrespective of this flag.
+             
+            momentum (float): The EMA coefficient used if track_running_stats is set to True. Higher value gives more weight to the current statistics (mean & std).
+            
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+
+        with torch.inference_mode(False):
+            model = nn.InstanceNorm2d(
+                num_features=num_features,
+                affine=affine,
+                track_running_stats=track_running_stats,
+                momentum=momentum)
+                    
+        return (model,)

modules/pytorch_wrapper/ptn_layer_norm.py

@@ -0,0 +1,82 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+from .utils import str_to_dim
+
+class PtnLayerNorm:
+    """
+    A normalization model to normalize elements over specified axes.
+
+        Args:
+            normalized_shape (list): Specifies the shape of elements to normalize.  
+                - For an input of shape `[8, 1024, 768]` (where axes represent `[batch_size, sequence_length, hidden_dim]`):
+                - `normalize_shape=[768]` normalizes within each token (across the hidden dimension).
+                - `normalize_shape=[1024, 768]` normalizes within each sample (across both sequence and token dimensions).  
+                - For an image input of shape `[8, 4, 256, 256]` (where axes represent `[batch_size, channel, height, width]`):
+                - `normalize_shape=[256, 256]` normalizes within each channel (across spatial dimensions).
+                - `normalize_shape=[4, 256, 256]` normalizes within each sample (across channels and spatial dimensions).
+
+            elementwise_affine (bool): If `True`, applies a learnable scaling factor to the normalized elements.
+
+            bias (bool): If `True` and elementwise_affine is also `True`, applies a learnable bias to the normalized elements.
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "normalized_shape": ("STRING", {"default": "[Specify shape here]"}),
+                "elementwise_affine": ("BOOLEAN", {"default": True}),
+                "bias": ("BOOLEAN", {"default": True}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          normalized_shape: str,
+          elementwise_affine: bool,
+          bias: bool
+          ) -> Tuple[nn.Module]:
+        """
+        A normalization model to normalize elements over specified axes.
+
+        Args:
+            normalized_shape (list): Specifies the shape of elements to normalize.  
+                - For an input of shape `[8, 1024, 768]` (where axes represent `[batch_size, sequence_length, hidden_dim]`):
+                - `normalize_shape=[768]` normalizes within each token (across the hidden dimension).
+                - `normalize_shape=[1024, 768]` normalizes within each sample (across both sequence and token dimensions).  
+                - For an image input of shape `[8, 4, 256, 256]` (where axes represent `[batch_size, channel, height, width]`):
+                - `normalize_shape=[256, 256]` normalizes within each channel (across spatial dimensions).
+                - `normalize_shape=[4, 256, 256]` normalizes within each sample (across channels and spatial dimensions).
+
+            elementwise_affine (bool): If `True`, applies a learnable scaling factor to the normalized elements.
+
+            bias (bool): If `True` and elementwise_affine is also `True`, applies a learnable bias to the normalized elements.
+            
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        shape = str_to_dim(normalized_shape)
+
+        with torch.inference_mode(False):
+            model = nn.LayerNorm(
+                normalized_shape=shape,
+                elementwise_affine=elementwise_affine,
+                bias=bias)
+            
+        return (model,)

modules/pytorch_wrapper/ptn_max_pool2d.py

@@ -5,6 +5,7 @@ from .utils import str_to_dim
 
 class PtnMaxPool2d:
     """
+    Ptn Max Pool 2d:
     A maxpool layer.  
 
         Args:

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "pt-wrapper"
 description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training."
-version = "3.1.3"
+version = "3.1.4"
 license = {file = "LICENSE"}
 
 [project.urls]

tests/pytorch_wrapper/ptn_batch_norm2d_test.py

@@ -0,0 +1,64 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_batch_norm2d import PtnBatchNorm2d
+
+
+class TestPtnBatchNorm2d(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnBatchNorm2d()
+
+    def check_normalization(self, x, num_features):
+        model = self.model_node.f(num_features,
+                                  affine=False,
+                                  track_running_stats=False,
+                                  momentum=0.001)[0]
+        
+        out = model(x)
+        self.assertEqual(out.shape, x.shape)
+
+        # Compute std along normalized axes before and after
+        std_before = x.std(dim=(0, 2, 3), unbiased=True)
+        std_after = out.std(dim=(0, 2, 3), unbiased=True)
+        mean_after = out.mean(dim=(0, 2, 3))
+
+        # Ensure std is not one before normalization
+        self.assertFalse(
+            torch.allclose(
+                std_before,
+                torch.ones_like(std_after),
+                atol=1e-1))
+
+        # Ensure normalization happened (std should be 1 unless elements are very close)
+        self.assertTrue(
+            torch.allclose(
+                std_after,
+                torch.ones_like(std_after),
+                atol=1e-1))
+    
+        self.assertTrue(
+            torch.allclose(
+                mean_after,
+                torch.zeros_like(mean_after),
+                atol=1e-1))
+
+    def test_normalization_cases(self):
+        test_cases = [
+            # Data, num_features on the channel axis
+            (torch.randn(2, 3, 3, 3) * 2, 3),  # c = 3
+            (torch.randn(3, 3, 3, 3) * 2, 3),  # c = 3
+        ]
+        
+        for x, norm_shape in test_cases:
+            with self.subTest(num_features=norm_shape):
+                self.check_normalization(x, norm_shape)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/ptn_instance_norm2d_test.py

@@ -0,0 +1,63 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_instance_norm2d import PtnInstanceNorm2d
+
+
+class TestPtnInstanceNorm2d(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnInstanceNorm2d()
+
+    def check_normalization(self, x, num_features):
+        model = self.model_node.f(num_features, True, True, 0.001)[0]
+        out = model(x)
+        self.assertEqual(out.shape, x.shape)
+        
+        # Compute std along normalized axes before and after
+        std_before = x.std(dim=(-2,-1), unbiased=False)
+        std_after = out.std(dim=(-2,-1), unbiased=False)
+        
+        mean_after = out.mean(dim=(-2,-1))
+
+        # Ensure std is not one before normalization
+        self.assertFalse(
+            torch.allclose(
+                std_before,
+                torch.ones_like(std_after),
+                atol=1e-3))
+
+        # Ensure normalization happened (std should be 1 unless elements are very close)
+        self.assertTrue(
+            torch.allclose(
+                std_after,
+                torch.ones_like(std_after),
+                atol=1e-3))
+    
+        self.assertTrue(
+            torch.allclose(
+                mean_after,
+                torch.zeros_like(mean_after),
+                atol=1e-3))
+
+    def test_normalization_cases(self):
+        test_cases = [
+            # Data, num_features on the channel axis
+            (torch.randn(3, 2, 2) * 2, 3),  # c = 3
+            (torch.randn(1, 3, 2, 2) * 2, 3),  # c = 3
+            (torch.randn(2, 3, 2, 2) * 2, 3),  # c = 3
+            (torch.randn(8, 3, 2, 2) * 2, 3),  # c = 3
+        ]
+        
+        for x, norm_shape in test_cases:
+            with self.subTest(num_features=norm_shape):
+                self.check_normalization(x, norm_shape)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/ptn_layer_norm_test.py

@@ -0,0 +1,70 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_layer_norm import PtnLayerNorm
+
+
+class TestPtnLayerNorm(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnLayerNorm()
+
+    def check_normalization(self, x, normalized_shape):
+        model = self.model_node.f(str(normalized_shape), True, True)[0]
+        out = model(x)
+        self.assertEqual(out.shape, x.shape)
+        
+        # Compute std along normalized axes before and after
+        std_before = x.std(dim=tuple(
+            range(-len(normalized_shape), 0)
+            ), unbiased=False)
+        std_after = out.std(dim=tuple(
+            range(-len(normalized_shape), 0)  # (-2, -1) for last two or (-3, -2, -1) for last three
+            ), unbiased=False)
+        
+        mean_after = out.mean(dim=tuple(
+            range(-len(normalized_shape), 0)  # (-2, -1) for last two or (-3, -2, -1) for last three
+            ))
+
+        # Ensure std is not one before normalization
+        self.assertFalse(
+            torch.allclose(
+                std_before,
+                torch.ones_like(std_after),
+                atol=1e-3))
+
+        # Ensure normalization happened (std should be 1 unless elements are very close)
+        self.assertTrue(
+            torch.allclose(
+                std_after,
+                torch.ones_like(std_after),
+                atol=1e-3))
+    
+        self.assertTrue(
+            torch.allclose(
+                mean_after,
+                torch.zeros_like(mean_after),
+                atol=1e-3))
+
+    def test_normalization_cases(self):
+        test_cases = [
+            # Data, normalized axis
+            (torch.randn(1, 2, 4) * 2, [4]),  # token
+            (torch.randn(1, 2, 4) * 2, [2, 4]),  # seq and token => per sample
+            (torch.randn(2, 2, 4) * 2, [2, 4]),  # seq and token => per sample
+            (torch.randn(1, 3, 2, 2) * 2, [2, 2]), # h, w
+            (torch.randn(1, 3, 2, 2) * 2, [3, 2, 2]), # c, h, w => per sample
+        ]
+        
+        for x, norm_shape in test_cases:
+            with self.subTest(normalized_shape=norm_shape):
+                self.check_normalization(x, norm_shape)
+
+
+if __name__ == "__main__":
+    unittest.main()