Merge branch 'main' into add-autojac-jac

Author: ValerianRey

.github/workflows/check-todos.yml

@@ -0,0 +1,34 @@
+name: "Check for TODOs"
+
+on:
+  pull_request:
+
+jobs:
+  check-todos:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Scan for TODO strings
+        run: |
+          echo "Scanning codebase for TODOs..."
+
+          git grep -nE "TODO" -- . ':(exclude).github/workflows/*' > todos_found.txt || true
+
+          if [ -s todos_found.txt ]; then
+            echo "❌ ERROR: Found TODOs in the following files:"
+            echo "-------------------------------------------"
+
+            while IFS=: read -r file line content; do
+              echo "::error file=$file,line=$line::TODO found at $file:$line - must be resolved before merge:%0A$content"
+            done < todos_found.txt
+
+            echo "-------------------------------------------"
+            echo "Please resolve these TODOs or track them in an issue before merging."
+
+            exit 1
+          else
+            echo "✅ No TODOs found. Codebase is clean!"
+            exit 0
+          fi

.github/workflows/claude.yml

@@ -0,0 +1,49 @@
+name: Claude Code
+
+on:
+  issue_comment:
+    types: [created]
+  pull_request_review_comment:
+    types: [created]
+  issues:
+    types: [opened, assigned]
+  pull_request_review:
+    types: [submitted]
+
+jobs:
+  claude:
+    if: |
+      (github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
+      (github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
+      (github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) ||
+      (github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude')))
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: read
+      issues: read
+      id-token: write
+      actions: read # Required for Claude to read CI results on PRs
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 1
+
+      - name: Run Claude Code
+        id: claude
+        uses: anthropics/claude-code-action@v1
+        with:
+          claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+
+          # This is an optional setting that allows Claude to read CI results on PRs
+          additional_permissions: |
+            actions: read
+
+          # Optional: Give a custom prompt to Claude. If this is not specified, Claude will perform the instructions specified in the comment that tagged it.
+          # prompt: 'Update the pull request description to include a summary of changes.'
+
+          # Optional: Add claude_args to customize behavior and configuration
+          # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
+          # or https://code.claude.com/docs/en/cli-reference for available options
+          # claude_args: '--allowed-tools Bash(gh pr:*)'

.github/workflows/tests.yml

@@ -52,8 +52,7 @@ jobs:
         env:
           PYTEST_TORCH_DTYPE: ${{ matrix.dtype || 'float32' }}
 
-      - &upload-codecov
-        name: Upload results to Codecov
+      - name: Upload results to Codecov
         uses: codecov/codecov-action@v4
         with:
           token: ${{ secrets.CODECOV_TOKEN }}

.gitignore

@@ -1,3 +1,6 @@
+# Profiling results
+traces/
+
 # uv
 uv.lock
 

.pre-commit-config.yaml

@@ -18,7 +18,7 @@ repos:
         ]
 
 -   repo: https://github.com/pycqa/isort
-    rev: 6.1.0
+    rev: 7.0.0
     hooks:
     -   id: isort  # Sort imports.
         args: [
@@ -31,7 +31,7 @@ repos:
         ]
 
 -   repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 25.9.0
+    rev: 25.12.0
     hooks:
     -   id: black  # Format code.
         args: [--line-length=100]

CHANGELOG.md

@@ -3,8 +3,8 @@
 All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
-and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). This changelog does not include internal
-changes that do not affect the user.
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). This
+changelog does not include internal changes that do not affect the user.
 
 ## [Unreleased]
 
@@ -13,7 +13,47 @@ changes that do not affect the user.
 - Added the function `torchjd.autojac.jac` to compute the Jacobian of some outputs with respect to
   some inputs, without doing any aggregation. Its interface is very similar to
   `torch.autograd.grad`.
-- Added `__all__` in the `__init__.py` of packages. This should prevent PyLance from triggering warnings when importing from `torchjd`.
+- Added a `scale_mode` parameter to `AlignedMTL` and `AlignedMTLWeighting`, allowing to choose
+  between `"min"`, `"median"`, and `"rmse"` scaling.
+
+### Changed
+
+- **BREAKING**: Removed from `backward` and `mtl_backward` the responsibility to aggregate the
+  Jacobian. Now, these functions compute and populate the `.jac` fields of the parameters, and a new
+  function `torchjd.autojac.jac_to_grad` should then be called to aggregate those `.jac` fields into
+  `.grad` fields.
+  This means that users now have more control on what they do with the Jacobians (they can easily
+  aggregate them group by group or even param by param if they want), but it now requires an extra
+  line of code to do the Jacobian descent step. To update, please change:
+  ```python
+  backward(losses, aggregator)
+  ```
+  to
+  ```python
+  backward(losses)
+  jac_to_grad(model.parameters(), aggregator)
+  ```
+  and
+  ```python
+  mtl_backward(losses, features, aggregator)
+  ```
+  to
+  ```python
+  mtl_backward(losses, features)
+  jac_to_grad(shared_module.parameters(), aggregator)
+  ```
+
+- Removed an unnecessary memory duplication. This should significantly improve the memory efficiency
+  of `autojac`.
+- Removed an unnecessary internal cloning of gradient. This should slightly improve the memory
+  efficiency of `autojac`.
+
+## [0.8.1] - 2026-01-07
+
+### Added
+
+- Added `__all__` in the `__init__.py` of packages. This should prevent PyLance from triggering
+  warnings when importing from `torchjd`.
 
 ## [0.8.0] - 2025-11-13
 

README.md

@@ -111,11 +111,11 @@ using [UPGrad](https://torchjd.org/stable/docs/aggregation/upgrad/).
       loss1 = loss_fn(output1, target1)
       loss2 = loss_fn(output2, target2)
 
-      optimizer.zero_grad()
 -     loss = loss1 + loss2
 -     loss.backward()
 +     mtl_backward(losses=[loss1, loss2], features=features, aggregator=aggregator)
       optimizer.step()
+      optimizer.zero_grad()
 ```
 
 > [!NOTE]
@@ -150,12 +150,12 @@ Jacobian descent using [UPGrad](https://torchjd.org/stable/docs/aggregation/upgr
 -     loss = loss_fn(output, target)  # shape [1]
 +     losses = loss_fn(output, target)  # shape [16]
 
-      optimizer.zero_grad()
 -     loss.backward()
 +     gramian = engine.compute_gramian(losses)  # shape: [16, 16]
 +     weights = weighting(gramian)  # shape: [16]
 +     losses.backward(weights)
       optimizer.step()
+      optimizer.zero_grad()
 ```
 
 Lastly, you can even combine the two approaches by considering multiple tasks and each element of
@@ -201,10 +201,10 @@ for input, target1, target2 in zip(inputs, task1_targets, task2_targets):
     # Obtain the weights that lead to no conflict between reweighted gradients
     weights = weighting(gramian)  # shape: [16, 2]
 
-    optimizer.zero_grad()
     # Do the standard backward pass, but weighted using the obtained weights
     losses.backward(weights)
     optimizer.step()
+    optimizer.zero_grad()
 ```
 
 > [!NOTE]

docs/source/docs/autojac/index.rst

@@ -11,3 +11,4 @@ autojac
     backward.rst
     mtl_backward.rst
     jac.rst
+    jac_to_grad.rst

docs/source/docs/autojac/jac_to_grad.rst

@@ -0,0 +1,6 @@
+:hide-toc:
+
+jac_to_grad
+===========
+
+.. autofunction:: torchjd.autojac.jac_to_grad

docs/source/examples/amp.rst

@@ -12,15 +12,15 @@ case, the losses) should preferably be scaled with a `GradScaler
 following example shows the resulting code for a multi-task learning use-case.
 
 .. code-block:: python
-    :emphasize-lines: 2, 17, 27, 34, 36-38
+    :emphasize-lines: 2, 17, 27, 34-35, 37-38
 
     import torch
     from torch.amp import GradScaler
     from torch.nn import Linear, MSELoss, ReLU, Sequential
     from torch.optim import SGD
 
     from torchjd.aggregation import UPGrad
-    from torchjd.autojac import mtl_backward
+    from torchjd.autojac import mtl_backward, jac_to_grad
 
     shared_module = Sequential(Linear(10, 5), ReLU(), Linear(5, 3), ReLU())
     task1_module = Linear(3, 1)
@@ -48,10 +48,11 @@ following example shows the resulting code for a multi-task learning use-case.
             loss2 = loss_fn(output2, target2)
 
         scaled_losses = scaler.scale([loss1, loss2])
-        optimizer.zero_grad()
-        mtl_backward(losses=scaled_losses, features=features, aggregator=aggregator)
+        mtl_backward(losses=scaled_losses, features=features)
+        jac_to_grad(shared_module.parameters(), aggregator)
         scaler.step(optimizer)
         scaler.update()
+        optimizer.zero_grad()
 
 .. hint::
     Within the ``torch.autocast`` context, some operations may be done in ``float16`` type. For