Fix documentation

Author: ValerianRey

src/torchjd/aggregation/_constant.py

@@ -8,13 +8,6 @@
 
 
 class _ConstantWeighting(Weighting[None]):
-    """
-    :class:`~torchjd.aggregation._weighting_bases.Weighting` that returns constant, pre-determined
-    weights.
-
-    :param weights: The weights to return at each call.
-    """
-
     def __init__(self, weights: Tensor) -> None:
         if weights.dim() != 1:
             raise ValueError(
@@ -30,6 +23,13 @@ def forward(self, _: None, /) -> Tensor:
 
 
 class ConstantWeighting(FromNothingWeighting):
+    """
+    :class:`~torchjd.aggregation._weighting_bases.Weighting` that returns constant, pre-determined
+    weights.
+
+    :param weights: The weights to return at each call.
+    """
+
     def __init__(self, weights: Tensor) -> None:
         super().__init__(_ConstantWeighting(weights))
 

src/torchjd/aggregation/_mean.py

@@ -9,12 +9,6 @@
 
 
 class _MeanWeighting(Weighting[Structure]):
-    r"""
-    :class:`~torchjd.aggregation._weighting_bases.Weighting` that gives the weights
-    :math:`\begin{bmatrix} \frac{1}{m} & \dots & \frac{1}{m} \end{bmatrix}^T \in
-    \mathbb{R}^m`.
-    """
-
     def forward(self, structure: Structure, /) -> Tensor:
         device = structure.device
         dtype = structure.dtype
@@ -24,6 +18,12 @@ def forward(self, structure: Structure, /) -> Tensor:
 
 
 class MeanWeighting(FromStructureWeighting):
+    r"""
+    :class:`~torchjd.aggregation._weighting_bases.Weighting` that gives the weights
+    :math:`\begin{bmatrix} \frac{1}{m} & \dots & \frac{1}{m} \end{bmatrix}^T \in
+    \mathbb{R}^m`.
+    """
+
     def __init__(self) -> None:
         super().__init__(_MeanWeighting())
 

src/torchjd/aggregation/_random.py

@@ -10,18 +10,18 @@
 
 
 class _RandomWeighting(Weighting[Structure]):
-    """
-    :class:`~torchjd.aggregation._weighting_bases.Weighting` that generates positive random weights
-    at each call.
-    """
-
     def forward(self, structure: Structure, /) -> Tensor:
         random_vector = torch.randn(structure.m, device=structure.device, dtype=structure.dtype)
         weights = F.softmax(random_vector, dim=-1)
         return weights
 
 
 class RandomWeighting(FromStructureWeighting):
+    """
+    :class:`~torchjd.aggregation._weighting_bases.Weighting` that generates positive random weights
+    at each call.
+    """
+
     def __init__(self) -> None:
         super().__init__(_RandomWeighting())
 

src/torchjd/aggregation/_sum.py

@@ -15,6 +15,11 @@ def forward(self, structure: Structure, /) -> Tensor:
 
 
 class SumWeighting(FromStructureWeighting):
+    r"""
+    :class:`~torchjd.aggregation._weighting_bases.Weighting` that gives the weights
+    :math:`\begin{bmatrix} 1 & \dots & 1 \end{bmatrix}^T \in \mathbb{R}^m`.
+    """
+
     def __init__(self) -> None:
         super().__init__(_SumWeighting())