ci: Test dependency lower bounds (#553)

* Increase permutation invariance tolerance for UPGrad
* Stop using nn.Buffer in tests (did not exist in torch 2.3.0)
* Add testing of dependency lower-bounds
* Increase lower bound of numpy to 1.21.2 (from 1.21.0) and of torch to 2.3.0 (from 2.0.0)
* Add changelog entry

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

CHANGELOG.md

@@ -50,6 +50,10 @@ changelog does not include internal changes that do not affect the user.
   of `autojac`.
 - Removed an unnecessary internal cloning of gradient. This should slightly improve the memory
   efficiency 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
 

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 = [
@@ -83,7 +83,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 +92,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 = [

tests/unit/aggregation/test_upgrad.py

@@ -33,7 +33,7 @@ def test_non_conflicting(aggregator: UPGrad, matrix: Tensor):
 
 @mark.parametrize(["aggregator", "matrix"], typical_pairs)
 def test_permutation_invariant(aggregator: UPGrad, matrix: Tensor):
-    assert_permutation_invariant(aggregator, matrix, n_runs=5, atol=4e-07, rtol=4e-07)
+    assert_permutation_invariant(aggregator, matrix, n_runs=5, atol=5e-07, rtol=5e-07)
 
 
 @mark.parametrize(["aggregator", "matrix"], typical_pairs)

tests/utils/architectures.py

@@ -595,9 +595,11 @@ class WithBuffered(ShapedModule):
     OUTPUT_SHAPES = (10,)
 
     class _Buffered(nn.Module):
+        buffer: Tensor
+
         def __init__(self):
             super().__init__()
-            self.buffer = nn.Buffer(torch.tensor(1.5))
+            self.register_buffer("buffer", torch.tensor(1.5))
 
         def forward(self, input: Tensor) -> Tensor:
             return input * self.buffer
@@ -636,7 +638,7 @@ class WithSideEffect(ShapedModule):
     def __init__(self):
         super().__init__()
         self.matrix = nn.Parameter(torch.randn(9, 10))
-        self.buffer = nn.Buffer(torch.zeros((9,)))
+        self.register_buffer("buffer", torch.zeros((9,)))
 
     def forward(self, input: Tensor) -> Tensor:
         self.buffer = self.buffer + 1.0