Maintain usage examples and their tests (#215)

* Use only floating point values in the definition of J in TrimmedMean's usage example
* Reorder imports in rst examples
* Remove filterwarnings in cagrad and nash_mtl usage examples
* Add module docstrings for tests/doc modules
* Wrap extra code of doc tests into Extra blocks

Author: ValerianRey

docs/source/examples/basic_usage.rst

@@ -15,7 +15,7 @@ Import several classes from ``torch`` and ``torchjd``:
 .. code-block:: python
 
     import torch
-    from torch.nn import MSELoss, Sequential, Linear, ReLU
+    from torch.nn import Linear, MSELoss, ReLU, Sequential
     from torch.optim import SGD
 
     import torchjd

docs/source/examples/iwrm.rst

@@ -29,10 +29,10 @@ each Jacobian matrix consists of one gradient per loss. In this example, we use
 
             import torch
             from torch.nn import (
-                MSELoss,
-                Sequential,
                 Linear,
-                ReLU
+                MSELoss,
+                ReLU,
+                Sequential
             )
             from torch.optim import SGD
 
@@ -70,10 +70,10 @@ each Jacobian matrix consists of one gradient per loss. In this example, we use
 
             import torch
             from torch.nn import (
-                MSELoss,
-                Sequential,
                 Linear,
-                ReLU
+                MSELoss,
+                ReLU,
+                Sequential
             )
             from torch.optim import SGD
 

src/torchjd/aggregation/cagrad.py

@@ -21,9 +21,6 @@ class CAGrad(_WeightedAggregator):
 
         Use CAGrad to aggregate a matrix.
 
-        >>> import warnings
-        >>> warnings.filterwarnings("ignore")
-        >>>
         >>> from torch import tensor
         >>> from torchjd.aggregation import CAGrad
         >>>

src/torchjd/aggregation/nash_mtl.py

@@ -51,9 +51,6 @@ class NashMTL(_WeightedAggregator):
 
         Use NashMTL to aggregate a matrix.
 
-        >>> import warnings
-        >>> warnings.filterwarnings("ignore")
-        >>>
         >>> from torch import tensor
         >>> from torchjd.aggregation import NashMTL
         >>>

src/torchjd/aggregation/trimmed_mean.py

@@ -25,10 +25,10 @@ class TrimmedMean(Aggregator):
         >>>
         >>> A = TrimmedMean(trim_number=1)
         >>> J = tensor([
-        ...     [ 1e11,     3],
-        ...     [    1, -1e11],
+        ...     [ 1e11,    3.],
+        ...     [   1., -1e11],
         ...     [-1e10,  1e10],
-        ...     [    2,     2],
+        ...     [   2.,    2.],
         ... ])
         >>>
         >>> A(J)

tests/doc/test_aggregation.py

@@ -1,3 +1,10 @@
+"""
+This file contains tests for the usage examples provided in the Aggregator subclasses. Each
+Aggregator's usage example, including its imports, should be copied in a function here, with the
+only difference that the advertised output should be replaced by a call to `assert_close`. The
+functions should be in alphabetical order.
+"""
+
 import torch
 from torch.testing import assert_close
 
@@ -14,9 +21,11 @@ def test_aligned_mtl():
 
 
 def test_cagrad():
+    # Extra ----------------------------------------------------------------------------------------
     import warnings
 
     warnings.filterwarnings("ignore")
+    # ----------------------------------------------------------------------------------------------
 
     from torch import tensor
 
@@ -116,9 +125,11 @@ def test_mgda():
 
 
 def test_nash_mtl():
+    # Extra ----------------------------------------------------------------------------------------
     import warnings
 
     warnings.filterwarnings("ignore")
+    # ----------------------------------------------------------------------------------------------
 
     from torch import tensor
 
@@ -149,7 +160,10 @@ def test_random():
     A = Random()
     J = tensor([[-4.0, 1.0, 1.0], [6.0, 1.0, 1.0]])
 
+    # Extra ----------------------------------------------------------------------------------------
     _ = torch.manual_seed(0)
+    # ----------------------------------------------------------------------------------------------
+
     assert_close(A(J), tensor([-2.6229, 1.0000, 1.0000]), rtol=0, atol=1e-4)
     assert_close(A(J), tensor([5.3976, 1.0000, 1.0000]), rtol=0, atol=1e-4)
 
@@ -173,10 +187,10 @@ def test_trimmed_mean():
     A = TrimmedMean(trim_number=1)
     J = tensor(
         [
-            [1e11, 3],
-            [1, -1e11],
+            [1e11, 3.0],
+            [1.0, -1e11],
             [-1e10, 1e10],
-            [2, 2],
+            [2.0, 2.0],
         ]
     )
 

tests/doc/test_backward.py

@@ -1,3 +1,8 @@
+"""
+This file contains the test of the backward usage example, with a verification of the value of the
+obtained `.grad` field.
+"""
+
 from torch.testing import assert_close
 from unit.conftest import DEVICE
 

tests/doc/test_rst.py

@@ -1,3 +1,10 @@
+"""
+This file contains the tests corresponding to the extra usage examples contained in the `.rst` files
+of the documentation. When there are multiple examples within a single `.rst` file, we use nested
+functions here to test them.
+"""
+
+
 def test_basic_usage():
     import torch
     from torch.nn import Linear, MSELoss, ReLU, Sequential
@@ -113,11 +120,13 @@ def test_mtl():
 
 
 def test_lightning_integration():
+    # Extra ----------------------------------------------------------------------------------------
     import logging
     import warnings
 
     warnings.filterwarnings("ignore")
     logging.disable(logging.INFO)
+    # ----------------------------------------------------------------------------------------------
 
     import torch
     from lightning import LightningModule, Trainer