Update

svekars · svekars · commit 477d0ff243c6 · 2025-05-16T10:10:57.000-07:00
diff --git a/torch/nn/modules/_nn_docs.py b/torch/nn/modules/_nn_docs.py
@@ -0,0 +1,5 @@
+# Common parameter documentation for nn modules
+COMMON_ARGS = """
+    device: the device on which the parameters will be allocated. Default: None
+    dtype: the data type of the parameters. Default: None
+"""
diff --git a/torch/nn/modules/normalization.py b/torch/nn/modules/normalization.py
@@ -8,6 +8,7 @@
 from torch.nn.parameter import Parameter
 
 from ._functions import CrossMapLRN2d as _cross_map_lrn2d
+from ._nn_docs import COMMON_ARGS
 from .module import Module
 
 
@@ -134,9 +135,6 @@ class LayerNorm(Module):
             and zeros (for biases). Default: ``True``.
         bias: If set to ``False``, the layer will not learn an additive bias (only relevant if
             :attr:`elementwise_affine` is ``True``). Default: ``True``.
-        device: the device on which the parameters will be allocated. Default: None
-        dtype: the data type of the parameters. Default: None
-
 
     Attributes:
         weight: the learnable weights of the module of shape
@@ -170,7 +168,9 @@ class LayerNorm(Module):
     .. image:: ../_static/img/nn/layer_norm.jpg
         :scale: 50 %
 
-    """
+    """.format(
+        COMMON_ARGS
+    )
 
     __constants__ = ["normalized_shape", "eps", "elementwise_affine"]
     normalized_shape: tuple[int, ...]
@@ -272,7 +272,9 @@ class GroupNorm(Module):
         >>> m = nn.GroupNorm(1, 6)
         >>> # Activating the module
         >>> output = m(input)
-    """
+    """.format(
+        COMMON_ARGS
+    )
 
     __constants__ = ["num_groups", "num_channels", "eps", "affine"]
     num_groups: int
@@ -349,8 +351,6 @@ class RMSNorm(Module):
         eps: a value added to the denominator for numerical stability. Default: :func:`torch.finfo(x.dtype).eps`
         elementwise_affine: a boolean value that when set to ``True``, this module
             has learnable per-element affine parameters initialized to ones (for weights). Default: ``True``.
-        device: the device on which the parameters will be allocated. Default: None.
-        dtype: the data type of the parameters. Default: None.
 
     Shape:
         - Input: :math:`(N, *)`
@@ -362,7 +362,9 @@ class RMSNorm(Module):
         >>> input = torch.randn(2, 2, 3)
         >>> rms_norm(input)
 
-    """
+    """.format(
+        COMMON_ARGS
+    )
 
     __constants__ = ["normalized_shape", "eps", "elementwise_affine"]
     normalized_shape: tuple[int, ...]
