Improve clarity of usage examples

Author: ValerianRey

README.md

@@ -149,8 +149,9 @@ Jacobian descent using [UPGrad](https://torchjd.org/stable/docs/aggregation/upgr
 
       optimizer.zero_grad()
 -     loss.backward()
-+     gramian = engine.compute_gramian(losses)
-+     losses.backward(weighting(gramian))
++     gramian = engine.compute_gramian(losses)  # shape: [16, 16]
++     weights = weighting(gramian)  # shape: [16]
++     losses.backward(weights)
       optimizer.step()
 ```
 

docs/source/examples/iwrm.rst

@@ -135,8 +135,8 @@ batch of data. When minimizing per-instance losses (IWRM), we use either autojac
                 y_hat = model(x).squeeze(dim=1)  # shape: [16]
                 losses = loss_fn(y_hat, y)  # shape: [16]
                 optimizer.zero_grad()
-                gramian = engine.compute_gramian(losses)
-                weights = weighting(gramian)
+                gramian = engine.compute_gramian(losses)  # shape: [16, 16]
+                weights = weighting(gramian)  # shape: [16]
                 losses.backward(weights)
                 optimizer.step()
 

src/torchjd/autogram/_engine.py

@@ -86,8 +86,9 @@ class Engine:
             >>>     losses = criterion(output, target)  # shape: [16]
             >>>
             >>>     optimizer.zero_grad()
-            >>>     gramian = engine.compute_gramian(losses)
-            >>>     losses.backward(weighting(gramian))
+            >>>     gramian = engine.compute_gramian(losses)  # shape: [16, 16]
+            >>>     weights = weighting(gramian)  # shape: [16]
+            >>>     losses.backward(weights)
             >>>     optimizer.step()
 
     .. warning::

tests/doc/test_autogram.py

@@ -25,6 +25,7 @@ def test_engine():
         losses = criterion(output, target)  # shape: [16]
 
         optimizer.zero_grad()
-        gramian = engine.compute_gramian(losses)
-        losses.backward(weighting(gramian))
+        gramian = engine.compute_gramian(losses)  # shape: [16, 16]
+        weights = weighting(gramian)  # shape: [16]
+        losses.backward(weights)
         optimizer.step()

tests/doc/test_rst.py

@@ -101,8 +101,9 @@ def test_autogram():
             y_hat = model(x).squeeze(dim=1)  # shape: [16]
             losses = loss_fn(y_hat, y)  # shape: [16]
             optimizer.zero_grad()
-            gramian = engine.compute_gramian(losses)
-            losses.backward(weighting(gramian))
+            gramian = engine.compute_gramian(losses)  # shape: [16, 16]
+            weights = weighting(gramian)  # shape: [16]
+            losses.backward(weights)
             optimizer.step()
 
     test_autograd()