feat: support reset_parameters()

build/torch26-cxx11-rocm62-x86_64-linux/activation/{_activation_d14fd4d_dirty.abi3.so → _activation_f3b99fb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:179bfe6bd5484e81b1d8fa6cc3e2596837946a17f0761b0bb2521fd162669046
+oid sha256:7bf9f4d85d15bc4869292e6a293ec53b7658cee61284457ea727c4be435062f7
 size 2656296

build/torch26-cxx11-rocm62-x86_64-linux/activation/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _activation_d14fd4d_dirty
-ops = torch.ops._activation_d14fd4d_dirty
+from . import _activation_f3b99fb_dirty
+ops = torch.ops._activation_f3b99fb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_activation_d14fd4d_dirty::{op_name}"
+    return f"_activation_f3b99fb_dirty::{op_name}"

build/torch26-cxx11-rocm62-x86_64-linux/activation/layers.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+from torch.nn import init
 
 from .poly_norm import PolyNormFunction
 from .rms_norm import RMSNormFunction
@@ -18,6 +19,13 @@ class PolyNorm(nn.Module):
     ):
         return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
 
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
 
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
@@ -30,3 +38,9 @@ class RMSNorm(nn.Module):
         x: torch.Tensor,
     ):
         return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)

build/torch27-cxx11-rocm63-x86_64-linux/activation/{_activation_d14fd4d_dirty.abi3.so → _activation_f3b99fb_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:94debfd52e15f782eb9dd328d9311080d803276745e440b176b20a7031299e3f
+oid sha256:ad289cf495aa7bcb7318535f2d76a6543bd44827369ec358ff7411e182ce089f
 size 2642736

build/torch27-cxx11-rocm63-x86_64-linux/activation/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _activation_d14fd4d_dirty
-ops = torch.ops._activation_d14fd4d_dirty
+from . import _activation_f3b99fb_dirty
+ops = torch.ops._activation_f3b99fb_dirty
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_activation_d14fd4d_dirty::{op_name}"
+    return f"_activation_f3b99fb_dirty::{op_name}"

build/torch27-cxx11-rocm63-x86_64-linux/activation/layers.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+from torch.nn import init
 
 from .poly_norm import PolyNormFunction
 from .rms_norm import RMSNormFunction
@@ -18,6 +19,13 @@ class PolyNorm(nn.Module):
     ):
         return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
 
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
 
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
@@ -30,3 +38,9 @@ class RMSNorm(nn.Module):
         x: torch.Tensor,
     ):
         return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)

torch-ext/activation/layers.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+from torch.nn import init
 
 from .poly_norm import PolyNormFunction
 from .rms_norm import RMSNormFunction
@@ -18,6 +19,13 @@ class PolyNorm(nn.Module):
     ):
         return PolyNormFunction.apply(x, self.weight, self.bias, self.eps)
 
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
 
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps=1e-6, dtype: torch.dtype = torch.float32):
@@ -30,3 +38,9 @@ class RMSNorm(nn.Module):
         x: torch.Tensor,
     ):
         return RMSNormFunction.apply(x, self.weight, self.eps)
+
+    def reset_parameters(self) -> None:
+        """
+        Resets parameters based on their initialization used in __init__.
+        """
+        init.ones_(self.weight)