Merge branch 'main' into add-autojac-jac

Author: ValerianRey

.github/actions/install-deps/action.yml

@@ -16,6 +16,10 @@ inputs:
 runs:
   using: composite
   steps:
+    - name: Create virtual environment
+      shell: bash
+      run: uv venv
+
     - name: Install dependencies (options=[${{ inputs.options }}], groups=[${{ inputs.groups }}])
       shell: bash
       run: |

.github/workflows/build-deploy-docs.yml

@@ -18,10 +18,10 @@ jobs:
       contents: write
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: '3.14'
 

.github/workflows/check-todos.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Scan for TODO strings
         run: |

.github/workflows/claude.yml

@@ -26,7 +26,7 @@ jobs:
       actions: read # Required for Claude to read CI results on PRs
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           fetch-depth: 1
 

.github/workflows/release.yml

@@ -14,10 +14,10 @@ jobs:
       id-token: write
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: '3.14'
 

.github/workflows/tests.yml

@@ -35,10 +35,10 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: ${{ matrix.python-version || '3.14' }}
 
@@ -62,10 +62,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: '3.14'
 
@@ -82,10 +82,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up uv
-        uses: astral-sh/setup-uv@v5
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: '3.14'
 

CHANGELOG.md

@@ -15,6 +15,9 @@ changelog does not include internal changes that do not affect the user.
   Its interface is analog to that of `torch.autograd.grad`.
 - Added a `scale_mode` parameter to `AlignedMTL` and `AlignedMTLWeighting`, allowing to choose
   between `"min"`, `"median"`, and `"rmse"` scaling.
+- Added an attribute `gramian_weighting` to all aggregators that use a gramian-based `Weighting`.
+  Usage is still the same, `aggregator.gramian_weighting` is just an alias for the (quite confusing)
+  `aggregator.weighting.weighting` field.
 
 ### Changed
 

docs/source/examples/monitoring.rst

@@ -49,7 +49,7 @@ they have a negative inner product).
     optimizer = SGD(params, lr=0.1)
     aggregator = UPGrad()
 
-    aggregator.weighting.weighting.register_forward_hook(print_weights)
+    aggregator.gramian_weighting.register_forward_hook(print_weights)
     aggregator.register_forward_hook(print_gd_similarity)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10

src/torchjd/aggregation/_aggregator_bases.py

@@ -73,8 +73,10 @@ class GramianWeightedAggregator(WeightedAggregator):
     WeightedAggregator that computes the gramian of the input jacobian matrix before applying a
     Weighting to it.
 
-    :param weighting: The object responsible for extracting the vector of weights from the gramian.
+    :param gramian_weighting: The object responsible for extracting the vector of weights from the
+        gramian.
     """
 
-    def __init__(self, weighting: Weighting[PSDMatrix]):
-        super().__init__(weighting << compute_gramian)
+    def __init__(self, gramian_weighting: Weighting[PSDMatrix]):
+        super().__init__(gramian_weighting << compute_gramian)
+        self.gramian_weighting = gramian_weighting

tests/doc/test_rst.py

@@ -308,7 +308,7 @@ def print_gd_similarity(_, inputs: tuple[torch.Tensor, ...], aggregation: torch.
     optimizer = SGD(params, lr=0.1)
     aggregator = UPGrad()
 
-    aggregator.weighting.weighting.register_forward_hook(print_weights)
+    aggregator.gramian_weighting.register_forward_hook(print_weights)
     aggregator.register_forward_hook(print_gd_similarity)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10