Add code/cube3d/model/transformers/norm.py

code/cube3d/model/transformers/norm.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+
+
+def fused_rms_norm(x: torch.Tensor, weight: nn.Parameter, eps: float):
+    """
+    Applies a fused Root Mean Square (RMS) normalization to the input tensor.
+    Args:
+        x (torch.Tensor): The input tensor to be normalized. Expected to have
+            at least one dimension.
+        weight (nn.Parameter): A learnable parameter used to scale the normalized
+            tensor. Its shape must be broadcastable to the shape of `x`.
+        eps (float): A small constant added to the denominator for numerical
+            stability during normalization.
+    Returns:
+        torch.Tensor: The normalized and scaled tensor with the same shape as `x`.
+    """
+
+    x = x.float()
+    return (x * torch.rsqrt((x * x).mean(-1, keepdim=True).add_(eps))) * weight
+
+
+class LayerNorm(nn.LayerNorm):
+    def forward(self, input: torch.Tensor):
+        """
+        Wrapper to ensure that the input tensor is cast to float before normalization.
+        """
+        y = super().forward(input.float())
+        return y.type_as(input)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine: bool = True):
+        """
+        Initializes the normalization layer.
+        Args:
+            dim (int): The number of features in the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Defaults to 1e-5.
+            elementwise_affine (bool, optional): If True, this layer will have learnable per-element affine parameters. Defaults to True.
+        """
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim), requires_grad=elementwise_affine)
+
+    def forward(self, x):
+        return fused_rms_norm(x, weight=self.weight, eps=self.eps).type_as(x)