Merge branch 'main' into optimize_jac_to_grad

Author: ValerianRey

.github/workflows/checks.yml

@@ -13,7 +13,7 @@ jobs:
   tests:
     # Default config: py3.14, ubuntu-latest, float32, full options.
     # The idea is to make each of those params vary one by one, to limit the number of tests to run.
-    name: Tests (py${{ matrix.python-version || '3.14' }}, ${{ matrix.os || 'ubuntu-latest' }}, ${{ matrix.dtype || 'float32' }}, ${{ matrix.options || 'full' }})
+    name: Tests (py${{ matrix.python-version || '3.14' }}, ${{ matrix.os || 'ubuntu-latest' }}, ${{ matrix.dtype || 'float32' }}, ${{ matrix.options || 'full' }}${{ matrix.extra_groups && format(', {0}', matrix.extra_groups) || '' }})
     runs-on: ${{ matrix.os || 'ubuntu-latest' }}
     strategy:
       fail-fast: false
@@ -32,6 +32,9 @@ jobs:
           - dtype: float64
           # Installation options variations
           - options: 'none'
+          # Lower-bounds of all dependencies and Python version.
+          - python-version: '3.10.0'
+            extra_groups: 'lower_bounds'
 
     steps:
       - name: Checkout repository
@@ -45,7 +48,7 @@ jobs:
       - uses: ./.github/actions/install-deps
         with:
           options: ${{ matrix.options || 'full' }}
-          groups: test
+          groups: test ${{ matrix.extra_groups }}
 
       - name: Run tests
         run: uv run pytest -W error tests/unit tests/doc --cov=src --cov-report=xml
@@ -106,8 +109,8 @@ jobs:
           # This reduces false positives due to rate limits
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
-  typing:
-    name: Typing correctness
+  code-quality:
+    name: Code quality (ty and ruff)
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
@@ -120,37 +123,10 @@ jobs:
 
       - uses: ./.github/actions/install-deps
         with:
-          groups: check
+          groups: check test plot
 
-      - name: Run mypy
-        run: uv run mypy src/torchjd
+      - name: Run ty
+        run: uv run ty check --output-format=github
 
-  check-todos:
-    name: Absence of TODOs
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v6
-
-      - 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
+      - name: Run ruff
+        run: uv run ruff check --output-format=github

.pre-commit-config.yaml

@@ -9,32 +9,12 @@ repos:
     -   id: check-docstring-first  # Check a common error of defining a docstring after code.
     -   id: check-merge-conflict  # Check for files that contain merge conflict strings.
 
--   repo: https://github.com/PyCQA/flake8
-    rev: 7.3.0
+-   repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.9.0
     hooks:
-    -   id: flake8  # Check style and syntax. Does not modify code, issues have to be solved manually.
-        args: [
-            '--ignore=E501,E203,W503,E402',  # Ignore line length problems, space after colon problems, line break occurring before a binary operator problems, module level import not at top of file problems.
-        ]
-
--   repo: https://github.com/pycqa/isort
-    rev: 7.0.0
-    hooks:
-    -   id: isort  # Sort imports.
-        args: [
-            --multi-line=3,
-            --line-length=100,
-            --trailing-comma,
-            --force-grid-wrap=0,
-            --use-parentheses,
-            --ensure-newline-before-comments,
-        ]
-
--   repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 25.12.0
-    hooks:
-    -   id: black  # Format code.
-        args: [--line-length=100]
+      - id: ruff
+        args: [ --fix, --ignore, FIX ]
+      - id: ruff-format
 
 ci:
     autoupdate_commit_msg: 'chore: Update pre-commit hooks'

CHANGELOG.md

@@ -45,9 +45,13 @@ changelog does not include internal changes that do not affect the user.
   mtl_backward(losses, features)
   jac_to_grad(shared_module.parameters(), aggregator)
   ```
-
+  
 - Removed several unnecessary memory duplications. This should significantly improve the memory
   efficiency and speed of `autojac`.
+- Increased the lower bounds of the torch (from 2.0.0 to 2.3.0) and numpy (from 1.21.0
+  to 1.21.2) dependencies to reflect what really works with torchjd. We now also run torchjd's tests
+  with the dependency lower-bounds specified in `pyproject.toml`, so we should now always accurately
+  reflect the actual lower-bounds.
 
 ## [0.8.1] - 2026-01-07
 

CONTRIBUTING.md

@@ -111,11 +111,12 @@ uv run pre-commit install
      uv run make clean
      ```
 
-## Running `mypy`
+## Type checking
 
-From the root of the repo, run:
+We use [ty](https://docs.astral.sh/ty/) for type-checking. If you're on VSCode, we recommend using
+the `ty` extension. You can also run it from the root of the repo with:
 ```bash
-uv run mypy src/torchjd
+uv run ty check
 ```
 
 ## Development guidelines

README.md

@@ -4,7 +4,6 @@
 [![Static Badge](https://img.shields.io/badge/%F0%9F%92%AC_ChatBot-chat.torchjd.org-blue?logo=%F0%9F%92%AC)](https://chat.torchjd.org)
 [![Tests](https://github.com/TorchJD/torchjd/actions/workflows/checks.yml/badge.svg)](https://github.com/TorchJD/torchjd/actions/workflows/checks.yml)
 [![codecov](https://codecov.io/gh/TorchJD/torchjd/graph/badge.svg?token=8AUCZE76QH)](https://codecov.io/gh/TorchJD/torchjd)
-[![mypy](https://img.shields.io/github/actions/workflow/status/TorchJD/torchjd/checks.yml?label=mypy)](https://github.com/TorchJD/torchjd/actions/workflows/checks.yml)
 [![pre-commit.ci status](https://results.pre-commit.ci/badge/github/TorchJD/torchjd/main.svg)](https://results.pre-commit.ci/latest/github/TorchJD/torchjd/main)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/torchjd)](https://pypi.org/project/torchjd/)
 [![Static Badge](https://img.shields.io/badge/Discord%20-%20community%20-%20%235865F2?logo=discord&logoColor=%23FFFFFF&label=Discord)](https://discord.gg/76KkRnb3nk)

pyproject.toml

@@ -13,9 +13,9 @@ authors = [
 ]
 requires-python = ">=3.10"
 dependencies = [
-    "torch>=2.0.0",
+    "torch>=2.3.0",  # Problems before 2.4.0, especially with autogram.
     "quadprog>=0.1.9, != 0.1.10",  # Doesn't work before 0.1.9, 0.1.10 is yanked
-    "numpy>=1.21.0",  # Does not work before 1.21
+    "numpy>=1.21.2",  # Does not work before 1.21. No python 3.10 wheel before 1.21.2.
     "qpsolvers>=1.0.1",  # Does not work before 1.0.1
 ]
 classifiers = [
@@ -66,7 +66,8 @@ Changelog = "https://github.com/TorchJD/torchjd/blob/main/CHANGELOG.md"
 
 [dependency-groups]
 check = [
-    "mypy>=1.16.0",
+    "ruff>=0.14.14",
+    "ty>=0.0.14",
     "pre-commit>=2.9.2",  # isort doesn't work before 2.9.2
 ]
 
@@ -83,7 +84,7 @@ test = [
     "pytest>=7.3",  # Before version 7.3, not all tests are run
     "pytest-cov>=6.0.0",  # Recent version to avoid problems, could be relaxed
     "lightning>=2.0.9",  # No OptimizerLRScheduler public type before 2.0.9
-    "torchvision>=0.22.1"  # Recent version to avoid problems, could be relaxed
+    "torchvision>=0.18.0"
 ]
 
 plot = [
@@ -92,6 +93,13 @@ plot = [
     "kaleido==0.2.1",  # Only works with locked version
     "matplotlib>=3.10.0",  # Recent version to avoid problems, could be relaxed
 ]
+# Dependency group allowing to easily resolve version of the core dependencies to the lower bound.
+lower_bounds = [
+    "torch==2.3.0",
+    "numpy==1.21.2",
+    "quadprog==0.1.9",
+    "qpsolvers==1.0.1",
+]
 
 [project.optional-dependencies]
 nash_mtl = [
@@ -114,3 +122,33 @@ exclude_lines = [
     "pragma: not covered",
     "@overload",
 ]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = [
+    "E",   # pycodestyle Error
+    "F",   # Pyflakes
+    "W",   # pycodestyle Warning
+    "I",   # isort
+    "UP",  # pyupgrade
+    "B",   # flake8-bugbear
+    "FIX",  # flake8-fixme
+]
+
+ignore = [
+    "E501", # line-too-long (handled by the formatter)
+    "E402", # module-import-not-at-top-of-file
+]
+
+[tool.ruff.lint.isort]
+combine-as-imports = true
+
+[tool.ruff.format]
+quote-style = "double"
+
+[tool.ty.src]
+include = ["src", "tests"]
+exclude = ["src/torchjd/aggregation/_nash_mtl.py"]

src/torchjd/__init__.py

@@ -1,8 +1,7 @@
 from collections.abc import Callable
 from warnings import warn as _warn
 
-from .autojac import backward as _backward
-from .autojac import mtl_backward as _mtl_backward
+from .autojac import backward as _backward, mtl_backward as _mtl_backward
 
 _deprecated_items: dict[str, tuple[str, Callable]] = {
     "backward": ("autojac", _backward),

src/torchjd/aggregation/_aligned_mtl.py

@@ -25,7 +25,7 @@
 # SOFTWARE.
 
 
-from typing import Literal
+from typing import Literal, TypeAlias
 
 import torch
 from torch import Tensor
@@ -37,6 +37,8 @@
 from ._utils.pref_vector import pref_vector_to_str_suffix, pref_vector_to_weighting
 from ._weighting_bases import Weighting
 
+SUPPORTED_SCALE_MODE: TypeAlias = Literal["min", "median", "rmse"]
+
 
 class AlignedMTL(GramianWeightedAggregator):
     r"""
@@ -58,10 +60,10 @@ class AlignedMTL(GramianWeightedAggregator):
     def __init__(
         self,
         pref_vector: Tensor | None = None,
-        scale_mode: Literal["min", "median", "rmse"] = "min",
+        scale_mode: SUPPORTED_SCALE_MODE = "min",
     ):
         self._pref_vector = pref_vector
-        self._scale_mode = scale_mode
+        self._scale_mode: SUPPORTED_SCALE_MODE = scale_mode
         super().__init__(AlignedMTLWeighting(pref_vector, scale_mode=scale_mode))
 
     def __repr__(self) -> str:
@@ -89,14 +91,14 @@ class AlignedMTLWeighting(Weighting[PSDMatrix]):
     def __init__(
         self,
         pref_vector: Tensor | None = None,
-        scale_mode: Literal["min", "median", "rmse"] = "min",
+        scale_mode: SUPPORTED_SCALE_MODE = "min",
     ):
         super().__init__()
         self._pref_vector = pref_vector
-        self._scale_mode = scale_mode
+        self._scale_mode: SUPPORTED_SCALE_MODE = scale_mode
         self.weighting = pref_vector_to_weighting(pref_vector, default=MeanWeighting())
 
-    def forward(self, gramian: PSDMatrix) -> Tensor:
+    def forward(self, gramian: PSDMatrix, /) -> Tensor:
         w = self.weighting(gramian)
         B = self._compute_balance_transformation(gramian, self._scale_mode)
         alpha = B @ w
@@ -105,7 +107,7 @@ def forward(self, gramian: PSDMatrix) -> Tensor:
 
     @staticmethod
     def _compute_balance_transformation(
-        M: Tensor, scale_mode: Literal["min", "median", "rmse"] = "min"
+        M: Tensor, scale_mode: SUPPORTED_SCALE_MODE = "min"
     ) -> Tensor:
         lambda_, V = torch.linalg.eigh(M, UPLO="U")  # More modern equivalent to torch.symeig
         tol = torch.max(lambda_) * len(M) * torch.finfo().eps

src/torchjd/aggregation/_cagrad.py

@@ -76,7 +76,7 @@ def __init__(self, c: float, norm_eps: float = 0.0001):
         self.c = c
         self.norm_eps = norm_eps
 
-    def forward(self, gramian: PSDMatrix) -> Tensor:
+    def forward(self, gramian: PSDMatrix, /) -> Tensor:
         U, S, _ = torch.svd(normalize(gramian, self.norm_eps))
 
         reduced_matrix = U @ S.sqrt().diag()

src/torchjd/aggregation/_constant.py

@@ -45,7 +45,7 @@ def __init__(self, weights: Tensor):
         super().__init__()
         self.weights = weights
 
-    def forward(self, matrix: Tensor) -> Tensor:
+    def forward(self, matrix: Tensor, /) -> Tensor:
         self._check_matrix_shape(matrix)
         return self.weights