Add testcode directives for doctests

Author: ValerianRey

docs/source/examples/amp.rst

@@ -11,7 +11,7 @@ case, the losses) should preferably be scaled with a `GradScaler
 <https://pytorch.org/docs/stable/amp.html#gradient-scaling>`_ to avoid gradient underflow. The
 following example shows the resulting code for a multi-task learning use-case.
 
-.. code-block:: python
+.. testcode::
     :emphasize-lines: 2, 17, 27, 34-35, 37-38
 
     import torch

docs/source/examples/basic_usage.rst

@@ -12,7 +12,7 @@ the parameters are updated using the resulting aggregation.
 
 Import several classes from ``torch`` and ``torchjd``:
 
-.. code-block:: python
+.. testcode::
 
     import torch
     from torch.nn import Linear, MSELoss, ReLU, Sequential
@@ -24,14 +24,14 @@ Import several classes from ``torch`` and ``torchjd``:
 
 Define the model and the optimizer, as usual:
 
-.. code-block:: python
+.. testcode::
 
     model = Sequential(Linear(10, 5), ReLU(), Linear(5, 2))
     optimizer = SGD(model.parameters(), lr=0.1)
 
 Define the aggregator that will be used to combine the Jacobian matrix:
 
-.. code-block:: python
+.. testcode::
 
     aggregator = UPGrad()
 
@@ -41,7 +41,7 @@ negatively affected by the update.
 
 Now that everything is defined, we can train the model. Define the input and the associated target:
 
-.. code-block:: python
+.. testcode::
 
     input = torch.randn(16, 10)  # Batch of 16 random input vectors of length 10
     target1 = torch.randn(16)  # First batch of 16 targets
@@ -51,7 +51,7 @@ Here, we generate fake inputs and labels for the sake of the example.
 
 We can now compute the losses associated to each element of the batch.
 
-.. code-block:: python
+.. testcode::
 
     loss_fn = MSELoss()
     output = model(input)
@@ -62,7 +62,7 @@ The last steps are similar to gradient descent-based optimization, but using the
 
 Perform the Jacobian descent backward pass:
 
-.. code-block:: python
+.. testcode::
 
     autojac.backward([loss1, loss2])
     jac_to_grad(model.parameters(), aggregator)
@@ -73,14 +73,14 @@ field of the parameters. It also deletes the ``.jac`` fields save some memory.
 
 Update each parameter based on its ``.grad`` field, using the ``optimizer``:
 
-.. code-block:: python
+.. testcode::
 
     optimizer.step()
 
 The model's parameters have been updated!
 
 As usual, you should now reset the ``.grad`` field of each model parameter:
 
-.. code-block:: python
+.. testcode::
 
     optimizer.zero_grad()

docs/source/examples/iwmtl.rst

@@ -9,7 +9,7 @@ this Gramian to reweight the gradients and resolve conflict entirely.
 
 The following example shows how to do that.
 
-.. code-block:: python
+.. testcode::
     :emphasize-lines: 5-6, 18-20, 31-32, 34-35, 37-38, 40-41
 
     import torch

docs/source/examples/iwrm.rst

@@ -41,7 +41,7 @@ batch of data. When minimizing per-instance losses (IWRM), we use either autojac
 .. tab-set::
     .. tab-item:: autograd (baseline)
 
-        .. code-block:: python
+        .. testcode::
 
             import torch
             from torch.nn import Linear, MSELoss, ReLU, Sequential
@@ -75,7 +75,7 @@ batch of data. When minimizing per-instance losses (IWRM), we use either autojac
 
     .. tab-item:: autojac
 
-        .. code-block:: python
+        .. testcode::
             :emphasize-lines: 5-6, 12, 16, 21-23
 
             import torch
@@ -110,7 +110,7 @@ batch of data. When minimizing per-instance losses (IWRM), we use either autojac
 
     .. tab-item:: autogram (recommended)
 
-        .. code-block:: python
+        .. testcode::
             :emphasize-lines: 5-6, 12, 16-17, 21-24
 
             import torch

docs/source/examples/lightning_integration.rst

@@ -10,7 +10,18 @@ The following code example demonstrates a basic multi-task learning setup using
 :class:`~lightning.pytorch.core.LightningModule` that will call :doc:`mtl_backward
 <../docs/autojac/mtl_backward>` at each training iteration.
 
-.. code-block:: python
+.. testsetup::
+
+    import warnings
+    import logging
+    from lightning.fabric.utilities.warnings import PossibleUserWarning
+
+    logging.disable(logging.INFO)
+    warnings.filterwarnings("ignore", category=DeprecationWarning)
+    warnings.filterwarnings("ignore", category=FutureWarning)
+    warnings.filterwarnings("ignore", category=PossibleUserWarning)
+
+.. testcode::
     :emphasize-lines: 9-10, 18, 31-32
 
     import torch

docs/source/examples/monitoring.rst

@@ -14,7 +14,12 @@ Jacobian descent is doing something different than gradient descent. With
 :doc:`UPGrad <../docs/aggregation/upgrad>`, this happens when the original gradients conflict (i.e.
 they have a negative inner product).
 
-.. code-block:: python
+.. testsetup::
+
+    import torch
+    torch.manual_seed(0)
+
+.. testcode::
     :emphasize-lines: 9-11, 13-18, 33-34
 
     import torch
@@ -67,3 +72,22 @@ they have a negative inner product).
         jac_to_grad(shared_module.parameters(), aggregator)
         optimizer.step()
         optimizer.zero_grad()
+
+.. testoutput::
+
+    Weights: tensor([0.5000, 0.5000])
+    Cosine similarity: 1.0000
+    Weights: tensor([0.5000, 0.5000])
+    Cosine similarity: 1.0000
+    Weights: tensor([0.5000, 0.5000])
+    Cosine similarity: 1.0000
+    Weights: tensor([0.6618, 1.0554])
+    Cosine similarity: 0.9249
+    Weights: tensor([0.6569, 1.2146])
+    Cosine similarity: 0.8661
+    Weights: tensor([0.5004, 0.5060])
+    Cosine similarity: 1.0000
+    Weights: tensor([0.5000, 0.5000])
+    Cosine similarity: 1.0000
+    Weights: tensor([0.5746, 1.1607])
+    Cosine similarity: 0.9301

docs/source/examples/mtl.rst

@@ -18,7 +18,7 @@ For the sake of the example, we generate a fake dataset consisting of 8 batches
 vectors of dimension 10, and their corresponding scalar labels for both tasks.
 
 
-.. code-block:: python
+.. testcode::
     :emphasize-lines: 5-6, 19, 32-33
 
     import torch

docs/source/examples/partial_jd.rst

@@ -13,7 +13,7 @@ perform the partial descent by considering only the parameters of the last two `
 doing this, we avoid computing the Jacobian and its Gramian with respect to the parameters of the
 first ``Linear`` layer, thereby reducing memory usage and computation time.
 
-.. code-block:: python
+.. testcode::
     :emphasize-lines: 16-18
 
     import torch

docs/source/examples/rnn.rst

@@ -5,7 +5,7 @@ When training recurrent neural networks for sequence modelling, we can easily ob
 element of the output sequences. If the gradients of these losses are likely to conflict, Jacobian
 descent can be leveraged to enhance optimization.
 
-.. code-block:: python
+.. testcode::
     :emphasize-lines: 5-6, 10, 17, 19-20
 
     import torch

src/torchjd/autogram/_engine.py

@@ -78,7 +78,7 @@ class Engine:
 
         Train a model using Gramian-based Jacobian descent.
 
-        .. code-block:: python
+        .. testcode::
             :emphasize-lines: 5-6, 15-16, 18-19, 26-29
 
             import torch
@@ -162,15 +162,21 @@ class Engine:
         Parent modules should call their child modules directly rather than using their child
         modules' parameters themselves. For instance, the following model is not supported:
 
-        >>> class Model(nn.Module):
-        >>>     def __init__(self):
-        >>>         super().__init__()
-        >>>         self.linear = nn.Linear(2, 3)  # Child module
-        >>>
-        >>>     def forward(self, input: Tensor) -> Tensor:
-        >>>         # Incorrect: Use the child module's parameters directly without calling it.
-        >>>         return input @ self.linear.weight.T + self.linear.bias
-        >>>         # Correct alternative: return self.linear(input)
+        .. testsetup::
+
+            from torch import nn
+
+        .. testcode::
+
+            class Model(nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.linear = nn.Linear(2, 3)  # Child module
+
+                def forward(self, input: Tensor) -> Tensor:
+                    # Incorrect: Use the child module's parameters directly without calling it.
+                    return input @ self.linear.weight.T + self.linear.bias
+                    # Correct alternative: return self.linear(input)
 
     .. note::
           For maximum efficiency, modules should ideally not contain both direct trainable