Make usage examples match doc tests

Author: ValerianRey

docs/source/examples/amp.rst

@@ -16,11 +16,11 @@ following example shows the resulting code for a multi-task learning use-case.
 
     import torch
     from torch.amp import GradScaler
-    from torch.nn import Sequential, Linear, ReLU, MSELoss
+    from torch.nn import Linear, MSELoss, ReLU, Sequential
     from torch.optim import SGD
 
-    from torchjd.autojac import mtl_backward
     from torchjd.aggregation import UPGrad
+    from torchjd.autojac import mtl_backward
 
     shared_module = Sequential(Linear(10, 5), ReLU(), Linear(5, 3), ReLU())
     task1_module = Linear(3, 1)

docs/source/examples/lightning_integration.rst

@@ -21,8 +21,8 @@ The following code example demonstrates a basic multi-task learning setup using
     from torch.optim import Adam
     from torch.utils.data import DataLoader, TensorDataset
 
-    from torchjd.autojac import mtl_backward
     from torchjd.aggregation import UPGrad
+    from torchjd.autojac import mtl_backward
 
     class Model(LightningModule):
         def __init__(self):

docs/source/examples/monitoring.rst

@@ -22,8 +22,8 @@ they have a negative inner product).
     from torch.nn.functional import cosine_similarity
     from torch.optim import SGD
 
-    from torchjd.autojac import mtl_backward
     from torchjd.aggregation import UPGrad
+    from torchjd.autojac import mtl_backward
 
     def print_weights(_, __, weights: torch.Tensor) -> None:
         """Prints the extracted weights."""

docs/source/examples/mtl.rst

@@ -25,8 +25,8 @@ vectors of dimension 10, and their corresponding scalar labels for both tasks.
     from torch.nn import Linear, MSELoss, ReLU, Sequential
     from torch.optim import SGD
 
-    from torchjd.autojac import mtl_backward
     from torchjd.aggregation import UPGrad
+    from torchjd.autojac import mtl_backward
 
     shared_module = Sequential(Linear(10, 5), ReLU(), Linear(5, 3), ReLU())
     task1_module = Linear(3, 1)

docs/source/examples/rnn.rst

@@ -12,8 +12,8 @@ descent can be leveraged to enhance optimization.
     from torch.nn import RNN
     from torch.optim import SGD
 
-    from torchjd.autojac import backward
     from torchjd.aggregation import UPGrad
+    from torchjd.autojac import backward
 
     rnn = RNN(input_size=10, hidden_size=20, num_layers=2)
     optimizer = SGD(rnn.parameters(), lr=0.1)

src/torchjd/autojac/_backward.py

@@ -41,8 +41,8 @@ def backward(
 
             >>> import torch
             >>>
-            >>> from torchjd.autojac import backward
             >>> from torchjd.aggregation import UPGrad
+            >>> from torchjd.autojac import backward
             >>>
             >>> param = torch.tensor([1., 2.], requires_grad=True)
             >>> # Compute arbitrary quantities that are function of param