[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot] · pre-commit-ci[bot] · commit a0d68e397c15 · 2025-06-12T03:53:51.000Z
for more information, see https://pre-commit.ci
diff --git a/docs/source/examples/sparse.rst b/docs/source/examples/sparse.rst
@@ -2,7 +2,7 @@ Quick example
 ==============================
 
 TorchJD now offers helpers that make working with sparse adjacency matrices
-transparent.  
+transparent.
 The key entry-point is :pyfunc:`torchjd.sparse.sparse_mm`,
 a vmap-aware autograd function that replaces the usual
 ``torch.sparse.mm`` inside Jacobian Descent pipelines.
diff --git a/src/torchjd/_autojac/__init__.py b/src/torchjd/_autojac/__init__.py
@@ -1,3 +1,4 @@
+from torchjd.sparse import sparse_mm
+
 from ._backward import backward
 from ._mtl_backward import mtl_backward
-from torchjd.sparse import sparse_mm
diff --git a/src/torchjd/aggregation/__init__.py b/src/torchjd/aggregation/__init__.py
@@ -19,7 +19,6 @@
 from ._sum import Sum
 from ._trimmed_mean import TrimmedMean
 from ._upgrad import UPGrad
-
 from ._utils.check_dependencies import (
     OptionalDepsNotInstalledError as _OptionalDepsNotInstalledError,
 )
diff --git a/src/torchjd/sparse/__init__.py b/src/torchjd/sparse/__init__.py
@@ -14,6 +14,6 @@
 
 __all__ = ["sparse_mm"]
 
-# feature flag 
+# feature flag
 if os.getenv("TORCHJD_DISABLE_SPARSE", "0") != "1":
     enable_seamless_sparse()
diff --git a/src/torchjd/sparse/_autograd.py b/src/torchjd/sparse/_autograd.py
@@ -21,8 +21,8 @@ def forward(A_like: torch.Tensor, X: torch.Tensor) -> torch.Tensor:  # noqa: D40
         if X.dim() == 3:  # (B, N, d)
             B, N, d = X.shape
             X2d = X.reshape(B * N, d).view(N, B * d)
-            Y2d = _orig_sparse_mm(A, X2d)             # pragma: no cover
-            return Y2d.view(N, B, d).permute(1, 0, 2) # pragma: no cover
+            Y2d = _orig_sparse_mm(A, X2d)  # pragma: no cover
+            return Y2d.view(N, B, d).permute(1, 0, 2)  # pragma: no cover
 
         return _orig_sparse_mm(A, X)
 
@@ -47,16 +47,16 @@ def backward(ctx, dY: torch.Tensor) -> Tuple[None, torch.Tensor]:
 
     @staticmethod
     def vmap(info, in_dims, A_unbatched, X_batched):  # noqa: D401
-        A = A_unbatched            # shared
-        X = X_batched              # (B, N, d)
+        A = A_unbatched  # shared
+        X = X_batched  # (B, N, d)
 
         B, N, d = X.shape
         X2d = X.reshape(B * N, d).view(N, B * d)
         Y2d = _orig_sparse_mm(A, X2d)
         Y = Y2d.view(N, B, d).permute(1, 0, 2)
-        return Y, 0                # output & out-dims
+        return Y, 0  # output & out-dims
 
 
 def sparse_mm(A_like: torch.Tensor, X: torch.Tensor) -> torch.Tensor:
     """Return ``A @ X`` through the vmap-safe sparse Function."""
-    return _SparseMatMul.apply(A_like, X)
+    return _SparseMatMul.apply(A_like, X)
diff --git a/src/torchjd/sparse/_patch.py b/src/torchjd/sparse/_patch.py
@@ -16,7 +16,6 @@
 
 from ._autograd import sparse_mm
 
-
 # The wheel might exist yet be ABI-incompatible with the current
 # PyTorch, which raises *OSError* at import-time.
 
@@ -29,6 +28,7 @@
 # Helpers
 def _wrap_mm(orig_fn: Callable, wrapper: Callable) -> Callable:
     """Return a patched ``torch.sparse.mm`` that defers to *wrapper*."""
+
     def patched(A, X):  # noqa: D401
         if isinstance(A, torch.Tensor) and A.is_sparse and X.dim() >= 2:
             return wrapper(A, X)
@@ -52,11 +52,9 @@ def enable_seamless_sparse() -> None:
     # torch.sparse.mm
     if getattr(torch.sparse, "_orig_mm", None) is None:
         torch.sparse._orig_mm = torch.sparse.mm  # type: ignore[attr-defined]
-        torch.sparse.mm = _wrap_mm(  # type: ignore[attr-defined]
-            torch.sparse._orig_mm, sparse_mm
-        )
+        torch.sparse.mm = _wrap_mm(torch.sparse._orig_mm, sparse_mm)  # type: ignore[attr-defined]
 
-    # tensor @ dense 
+    # tensor @ dense
     if getattr(torch.Tensor, "_orig_matmul", None) is None:
         torch.Tensor._orig_matmul = torch.Tensor.__matmul__  # type: ignore[attr-defined]  # noqa: E501
         torch.Tensor.__matmul__ = _wrap_tensor_matmul(
@@ -69,10 +67,11 @@ def enable_seamless_sparse() -> None:
             "torch_sparse not found: SpSpMM will use slow fallback.",
             RuntimeWarning,
             stacklevel=2,
-        ) # pragma: no cover
+        )  # pragma: no cover
         return
 
     if not hasattr(torch_sparse.SparseTensor, "_orig_matmul"):
+
         def _sparse_tensor_matmul(self, dense):  # noqa: D401
             return sparse_mm(self, dense)
 
diff --git a/src/torchjd/sparse/_utils.py b/src/torchjd/sparse/_utils.py
@@ -34,4 +34,4 @@ def to_coalesced_coo(x: Any) -> torch.Tensor:
     except ModuleNotFoundError:  # pragma: no cover
         pass
 
-    raise TypeError(f"Unsupported sparse type: {type(x)}") # pragma: no cover
+    raise TypeError(f"Unsupported sparse type: {type(x)}")  # pragma: no cover
diff --git a/tests/unit/sparse/test_mm.py b/tests/unit/sparse/test_mm.py
@@ -1,10 +1,13 @@
-import torch
 import pytest
+import torch
+
 from torchjd.sparse import sparse_mm
 from torchjd.sparse._utils import to_coalesced_coo
 
 try:
-    import importlib, types
+    import importlib
+    import types
+
     torch_sparse = importlib.import_module("torch_sparse")  # noqa: E402
     HAVE_TORCH_SPARSE = isinstance(torch_sparse, types.ModuleType)
 except (ModuleNotFoundError, OSError):
@@ -13,6 +16,7 @@
 
 try:
     import scipy.sparse as sp
+
     HAVE_SCIPY = True
 except ModuleNotFoundError:
     HAVE_SCIPY = False
@@ -32,7 +36,7 @@ def _batched_features(device):
 def test_vmap_branch(device):
     A = _dense_graph().to(device)
     X = _batched_features(device)
-    Y = sparse_mm(A, X)                # calls vmap-aware branch
+    Y = sparse_mm(A, X)  # calls vmap-aware branch
     assert Y.shape == X.shape
 
 
diff --git a/tests/unit/sparse/test_mm_3d.py b/tests/unit/sparse/test_mm_3d.py
@@ -1,16 +1,18 @@
 import torch
+
 from torchjd.sparse import sparse_mm
 
+
 def test_forward_backward_3d():
     # sparse 2×2 matrix
     A = torch.sparse_coo_tensor([[0, 1], [1, 0]], [1.0, 1.0]).coalesce()
 
     # 3-D dense tensor (B=3, N=2, d=4)
     X = torch.randn(3, 2, 4, requires_grad=True)
 
-    Y = sparse_mm(A, X)          # exercises 3-D forward branch
+    Y = sparse_mm(A, X)  # exercises 3-D forward branch
     loss = Y.sum()
-    loss.backward()              # exercises 3-D backward branch
+    loss.backward()  # exercises 3-D backward branch
 
     # Gradient should be ones because A.T @ 1 = [1,1] → broadcast
     assert torch.allclose(X.grad, torch.ones_like(X), atol=1e-6)
diff --git a/tests/unit/sparse/test_mm_sequential.py b/tests/unit/sparse/test_mm_sequential.py
@@ -1,8 +1,10 @@
 import torch
+
 from torchjd._autojac import backward
 from torchjd.aggregation import UPGrad
 from torchjd.sparse import sparse_mm
 
+
 def test_sequential_backward():
     A = torch.sparse_coo_tensor([[0, 1], [1, 0]], [1.0, 1.0]).coalesce()
     p = torch.tensor([1.0, 2.0], requires_grad=True)
diff --git a/tests/unit/sparse/test_mm_single.py b/tests/unit/sparse/test_mm_single.py
@@ -1,10 +1,12 @@
 import torch
+
 from torchjd.sparse import sparse_mm
 
+
 def test_single_forward_backward():
-    A = torch.sparse_coo_tensor([[0,1],[1,0]], [1.,1.]).coalesce()
+    A = torch.sparse_coo_tensor([[0, 1], [1, 0]], [1.0, 1.0]).coalesce()
     X = torch.randn(2, 5, requires_grad=True)
-    Y = sparse_mm(A, X)           # (2,5)
+    Y = sparse_mm(A, X)  # (2,5)
     loss = Y.sum()
     loss.backward()
     # gradient should equal A.T @ 1 = [1,1]
diff --git a/tests/unit/sparse/test_mm_vmap.py b/tests/unit/sparse/test_mm_vmap.py
@@ -1,13 +1,15 @@
 import torch
 from torch.func import vmap
+
 from torchjd.sparse import sparse_mm
 
+
 def test_batched_vmap_forward_backward():
     """
     Touch the custom vmap rule in _SparseMatMul to push per-file coverage
     above the 90 % guideline.
     """
-    A = torch.sparse_coo_tensor([[0, 1], [1, 0]], [1., 1.]).coalesce()
+    A = torch.sparse_coo_tensor([[0, 1], [1, 0]], [1.0, 1.0]).coalesce()
     B, N, d = 4, 2, 3
     X = torch.randn(B, N, d, requires_grad=True)
 
diff --git a/tests/unit/sparse/test_patch.py b/tests/unit/sparse/test_patch.py
@@ -1,10 +1,12 @@
 import torch
+
 from torchjd.sparse._patch import enable_seamless_sparse
 
+
 def test_monkey_patch_matmul():
     enable_seamless_sparse()  # idempotent
     A = torch.sparse_coo_tensor([[0, 1], [1, 0]], [1.0, 1.0]).coalesce()
     X = torch.randn(2, 3)
-    Y1 = A @ X                 # should hit sparse_mm via patched __matmul__
-    Y2 = torch.tensor([[0., 0., 0.], [0., 0., 0.]])  # placeholder
+    Y1 = A @ X  # should hit sparse_mm via patched __matmul__
+    Y2 = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])  # placeholder
     assert torch.allclose(Y1.sum(), (A.to_dense() @ X).sum())
diff --git a/tests/unit/sparse/test_patch_idempotent.py b/tests/unit/sparse/test_patch_idempotent.py
@@ -1,5 +1,6 @@
 from torchjd.sparse._patch import enable_seamless_sparse
 
+
 def test_enable_patch_idempotent():
-    enable_seamless_sparse()   # first call patches
-    enable_seamless_sparse()   # second call should be a no-op
+    enable_seamless_sparse()  # first call patches
+    enable_seamless_sparse()  # second call should be a no-op
diff --git a/tests/unit/sparse/test_patch_import.py b/tests/unit/sparse/test_patch_import.py
@@ -1,6 +1,9 @@
-import importlib, sys, types
+import importlib
+import sys
+import types
 from contextlib import contextmanager
 
+
 @contextmanager
 def fake_torch_sparse():
     """
@@ -15,7 +18,7 @@ class Dummy:  # noqa: D401
         def matmul(self, dense):
             raise NotImplementedError
 
-    mod.SparseTensor = Dummy                                    # type: ignore
+    mod.SparseTensor = Dummy  # type: ignore
     sys.modules["torch_sparse"] = mod
     try:
         yield
@@ -26,13 +29,18 @@ def matmul(self, dense):
 def test_patch_without_torch_sparse(monkeypatch):
     monkeypatch.setitem(sys.modules, "torch_sparse", None)
     from importlib import reload
+
     import torchjd.sparse._patch as p
-    reload(p)                      # re-import to trigger patch
+
+    reload(p)  # re-import to trigger patch
     assert p.torch_sparse is None  # slow fallback branch hit
 
+
 def test_patch_with_dummy_torch_sparse(monkeypatch):
     with fake_torch_sparse():
         from importlib import reload
+
         import torchjd.sparse._patch as p
+
         reload(p)
-        assert p.torch_sparse is not None   # optional branch hit
+        assert p.torch_sparse is not None  # optional branch hit
diff --git a/tests/unit/sparse/test_patch_torch_sparse_branch.py b/tests/unit/sparse/test_patch_torch_sparse_branch.py
@@ -1,6 +1,9 @@
-import importlib, sys, types
+import importlib
+import sys
+import types
 from importlib import reload
 
+
 def _make_dummy_torch_sparse():
     """
     Return a minimal torch_sparse stub:
@@ -16,6 +19,7 @@ def matmul(self, dense):
 
         def to_torch_sparse_coo_tensor(self):
             import torch
+
             return torch.sparse_coo_tensor([[0], [0]], [1.0], (1, 1))
 
     dummy_mod.SparseTensor = DummyTensor  # type: ignore[attr-defined]
@@ -27,10 +31,11 @@ def test_full_torch_sparse_branch(monkeypatch):
     monkeypatch.setitem(sys.modules, "torch_sparse", _make_dummy_torch_sparse())
 
     # Force the patch module to re-evaluate from scratch
-    import torchjd.sparse._patch as p  # noqa: E402
-
     # Remove earlier sentinel attributes so enable_seamless_sparse() re-patches
     import torch
+
+    import torchjd.sparse._patch as p  # noqa: E402
+
     for attr in ("_orig_mm",):
         if hasattr(torch.sparse, attr):
             delattr(torch.sparse, attr)  # type: ignore[attr-defined]
diff --git a/tests/unit/sparse/test_patch_warn_branch.py b/tests/unit/sparse/test_patch_warn_branch.py
@@ -3,7 +3,11 @@
 when *no* ``torch_sparse`` package is available.
 """
 
-import importlib, sys, types, warnings
+import importlib
+import sys
+import types
+import warnings
+
 import torch
 
 
@@ -14,12 +18,11 @@ def test_warn_branch(monkeypatch):
         delattr(torch.sparse, "_orig_mm")  # type: ignore[attr-defined]
 
     import torchjd.sparse._patch as p  # noqa: E402
+
     p = importlib.reload(p)
 
     with warnings.catch_warnings(record=True) as rec:
         warnings.simplefilter("always")
         p.enable_seamless_sparse()  # <- emits RuntimeWarning branch
 
-    assert any(
-        "SpSpMM will use slow fallback" in str(w.message) for w in rec
-    )
+    assert any("SpSpMM will use slow fallback" in str(w.message) for w in rec)
diff --git a/tests/unit/sparse/test_sparse_mm_wrapper.py b/tests/unit/sparse/test_sparse_mm_wrapper.py
@@ -1,11 +1,13 @@
 import torch
+
 from torchjd.sparse._patch import enable_seamless_sparse
 
+
 def test_torch_sparse_mm_wrapper():
-    enable_seamless_sparse()             # idempotent
-    A = torch.sparse_coo_tensor([[0, 1], [1, 0]], [1., 1.]).coalesce()
+    enable_seamless_sparse()  # idempotent
+    A = torch.sparse_coo_tensor([[0, 1], [1, 0]], [1.0, 1.0]).coalesce()
     X = torch.randn(2, 3)
 
-    out = torch.sparse.mm(A, X)          # routed through wrapper
+    out = torch.sparse.mm(A, X)  # routed through wrapper
     ref = A.to_dense() @ X
     assert torch.allclose(out, ref, atol=1e-6)
diff --git a/tests/unit/sparse/test_utils_scipy.py b/tests/unit/sparse/test_utils_scipy.py
@@ -1,10 +1,12 @@
 import importlib
-import pytest
+
 import numpy as np
+import pytest
 
 scipy = pytest.importorskip("scipy")  # skip if SciPy not available
 from torchjd.sparse._utils import to_coalesced_coo
 
+
 def test_to_coalesced_coo_from_scipy():
     sp = importlib.import_module("scipy.sparse")
     # 2×2 off-diagonal ones
diff --git a/tests/unit/sparse/test_utils_torch_sparse.py b/tests/unit/sparse/test_utils_torch_sparse.py
@@ -1,4 +1,9 @@
-import importlib, sys, types, torch
+import importlib
+import sys
+import types
+
+import torch
+
 
 def test_to_coalesced_coo_torch_sparse(monkeypatch):
     dummy = types.ModuleType("torch_sparse")
@@ -10,9 +15,7 @@ def __init__(self):
             self.value = torch.tensor([1.0])
 
         def to_torch_sparse_coo_tensor(self):
-            return torch.sparse_coo_tensor(
-                torch.stack([self.row, self.col]), self.value, (1, 1)
-            )
+            return torch.sparse_coo_tensor(torch.stack([self.row, self.col]), self.value, (1, 1))
 
         def matmul(self, other):
             raise NotImplementedError

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+from torchjd.sparse import sparse_mm`
	`2`	`+`
`1`	`3`	`from ._backward import backward`
`2`	`4`	`from ._mtl_backward import mtl_backward`
`3`		`-from torchjd.sparse import sparse_mm`
Original file line number	Diff line number	Diff line change
`@@ -19,7 +19,6 @@`
`19`	`19`	`from ._sum import Sum`
`20`	`20`	`from ._trimmed_mean import TrimmedMean`
`21`	`21`	`from ._upgrad import UPGrad`
`22`		`-`
`23`	`22`	`from ._utils.check_dependencies import (`
`24`	`23`	`OptionalDepsNotInstalledError as _OptionalDepsNotInstalledError,`
`25`	`24`	`)`