Import directly torchjd.autojac and use autojac.backward in basic usage example

Author: ValerianRey

docs/source/examples/basic_usage.rst

@@ -18,7 +18,7 @@ Import several classes from ``torch`` and ``torchjd``:
     from torch.nn import Linear, MSELoss, ReLU, Sequential
     from torch.optim import SGD
 
-    import torchjd
+    from torchjd import autojac
     from torchjd.aggregation import UPGrad
 
 Define the model and the optimizer, as usual:
@@ -69,7 +69,7 @@ Perform the Jacobian descent backward pass:
 
 .. code-block:: python
 
-    torchjd.autojac.backward([loss1, loss2], aggregator)
+    autojac.backward([loss1, loss2], aggregator)
 
 This will populate the ``.grad`` field of each model parameter with the corresponding aggregated
 Jacobian matrix.

tests/doc/test_rst.py

@@ -10,7 +10,7 @@ def test_basic_usage():
     from torch.nn import Linear, MSELoss, ReLU, Sequential
     from torch.optim import SGD
 
-    import torchjd
+    from torchjd import autojac
     from torchjd.aggregation import UPGrad
 
     model = Sequential(Linear(10, 5), ReLU(), Linear(5, 2))
@@ -27,7 +27,7 @@ def test_basic_usage():
     loss2 = loss_fn(output[:, 1], target2)
 
     optimizer.zero_grad()
-    torchjd.autojac.backward([loss1, loss2], aggregator)
+    autojac.backward([loss1, loss2], aggregator)
     optimizer.step()