typing: Replace mypy with ty (#551)

PierreQuinton · ValerianRey · web-flow · commit 65c1165f07ad · 2026-02-03T13:41:03.000+01:00
Typing changes (necessary to run ty without error):
* Use positional-only arguments for methods from classes whose subclasses rename the parameters. This includes `Weighting.forward`, `JacobianComputer._compute_jacobian`, `Transform.__call__` and `Differentiate.__differentiate__`. This fixes a break of LSP.
* Remove now-useless typing ignore statement in MGDA
* Add necessary casts in NashMTL
* Add ignore statement for subclasses of autograd.Functions
* Fix name of parameters of methods in `OrderedSet`
* Make ModuleFactory generic
* Add some missing casts to PSDMatrix
* Add ignore statement when calling .grad of BatchedTensor
* Ignore type errors in the lightning example's test

Structural changes:
* Add ty check dependency
* Remove mypy check dependency
* Change section about type checking in CONTRIBUTING.md
* Remove mypy badge
* Add ty tool section in pyproject.toml
* Change CI to run ty instead of mypy
* Check typing in tests too
---------
Co-authored-by: Valérian Rey &lt;valerian.rey@gmail.com&gt;
diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
@@ -120,10 +120,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
 
   check-todos:
     name: Absence of TODOs
diff --git a/CONTRIBUTING.md b/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
diff --git a/README.md b/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)
diff --git a/pyproject.toml b/pyproject.toml
@@ -66,7 +66,7 @@ Changelog = "https://github.com/TorchJD/torchjd/blob/main/CHANGELOG.md"
 
 [dependency-groups]
 check = [
-    "mypy>=1.16.0",
+    "ty>=0.0.14",
     "pre-commit>=2.9.2",  # isort doesn't work before 2.9.2
 ]
 
@@ -114,3 +114,7 @@ exclude_lines = [
     "pragma: not covered",
     "@overload",
 ]
+
+[tool.ty.src]
+include = ["src", "tests"]
+exclude = ["src/torchjd/aggregation/_nash_mtl.py"]
diff --git a/src/torchjd/aggregation/_aligned_mtl.py b/src/torchjd/aggregation/_aligned_mtl.py
@@ -98,7 +98,7 @@ def __init__(
         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
diff --git a/src/torchjd/aggregation/_cagrad.py b/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()
diff --git a/src/torchjd/aggregation/_constant.py b/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
 
diff --git a/src/torchjd/aggregation/_dualproj.py b/src/torchjd/aggregation/_dualproj.py
@@ -85,7 +85,7 @@ def __init__(
         self.reg_eps = reg_eps
         self.solver: SUPPORTED_SOLVER = solver
 
-    def forward(self, gramian: PSDMatrix) -> Tensor:
+    def forward(self, gramian: PSDMatrix, /) -> Tensor:
         u = self.weighting(gramian)
         G = regularize(normalize(gramian, self.norm_eps), self.reg_eps)
         w = project_weights(u, G, self.solver)
diff --git a/src/torchjd/aggregation/_imtl_g.py b/src/torchjd/aggregation/_imtl_g.py
@@ -29,7 +29,7 @@ class IMTLGWeighting(Weighting[PSDMatrix]):
     :class:`~torchjd.aggregation.IMTLG`.
     """
 
-    def forward(self, gramian: PSDMatrix) -> Tensor:
+    def forward(self, gramian: PSDMatrix, /) -> Tensor:
         d = torch.sqrt(torch.diagonal(gramian))
         v = torch.linalg.pinv(gramian) @ d
         v_sum = v.sum()
diff --git a/src/torchjd/aggregation/_krum.py b/src/torchjd/aggregation/_krum.py
@@ -61,7 +61,7 @@ def __init__(self, n_byzantine: int, n_selected: int = 1):
         self.n_byzantine = n_byzantine
         self.n_selected = n_selected
 
-    def forward(self, gramian: PSDMatrix) -> Tensor:
+    def forward(self, gramian: PSDMatrix, /) -> Tensor:
         self._check_matrix_shape(gramian)
         gradient_norms_squared = torch.diagonal(gramian)
         distances_squared = (
diff --git a/src/torchjd/aggregation/_mean.py b/src/torchjd/aggregation/_mean.py
@@ -24,7 +24,7 @@ class MeanWeighting(Weighting[Matrix]):
     \mathbb{R}^m`.
     """
 
-    def forward(self, matrix: Tensor) -> Tensor:
+    def forward(self, matrix: Tensor, /) -> Tensor:
         device = matrix.device
         dtype = matrix.dtype
         m = matrix.shape[0]
diff --git a/src/torchjd/aggregation/_mgda.py b/src/torchjd/aggregation/_mgda.py
@@ -43,7 +43,7 @@ def __init__(self, epsilon: float = 0.001, max_iters: int = 100):
         self.epsilon = epsilon
         self.max_iters = max_iters
 
-    def forward(self, gramian: PSDMatrix) -> Tensor:
+    def forward(self, gramian: PSDMatrix, /) -> Tensor:
         """
         This is the Frank-Wolfe solver in Algorithm 2 of `Multi-Task Learning as Multi-Objective
         Optimization
@@ -65,7 +65,7 @@ def forward(self, gramian: PSDMatrix) -> Tensor:
             elif b <= a:
                 gamma = 0.0
             else:
-                gamma = (b - a) / (b + c - 2 * a)  # type: ignore[assignment]
+                gamma = (b - a) / (b + c - 2 * a)
             alpha = (1 - gamma) * alpha + gamma * e_t
             if gamma < self.epsilon:
                 break
diff --git a/src/torchjd/aggregation/_nash_mtl.py b/src/torchjd/aggregation/_nash_mtl.py
@@ -23,7 +23,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-# mypy: ignore-errors
+from typing import cast
 
 from torchjd._linalg import Matrix
 
@@ -96,7 +96,7 @@ def __init__(
 
     def reset(self) -> None:
         """Resets the internal state of the algorithm."""
-        self.weighting.reset()
+        cast(_NashMTLWeighting, self.weighting).reset()
 
     def __repr__(self) -> str:
         return (
@@ -141,7 +141,7 @@ def __init__(
         self.prvs_alpha = np.ones(self.n_tasks, dtype=np.float32)
 
     def _stop_criteria(self, gtg: np.ndarray, alpha_t: np.ndarray) -> bool:
-        return (
+        return bool(
             (self.alpha_param.value is None)
             or (np.linalg.norm(gtg @ alpha_t - 1 / (alpha_t + 1e-10)) < 1e-3)
             or (np.linalg.norm(self.alpha_param.value - self.prvs_alpha_param.value) < 1e-6)
@@ -198,7 +198,7 @@ def _init_optim_problem(self) -> None:
         obj = cp.Minimize(cp.sum(G_alpha) + self.phi_alpha / self.normalization_factor_param)
         self.prob = cp.Problem(obj, constraint)
 
-    def forward(self, matrix: Tensor) -> Tensor:
+    def forward(self, matrix: Tensor, /) -> Tensor:
         if self.step == 0:
             self._init_optim_problem()
 
diff --git a/src/torchjd/aggregation/_pcgrad.py b/src/torchjd/aggregation/_pcgrad.py
@@ -29,7 +29,7 @@ class PCGradWeighting(Weighting[PSDMatrix]):
     :class:`~torchjd.aggregation.PCGrad`.
     """
 
-    def forward(self, gramian: PSDMatrix) -> Tensor:
+    def forward(self, gramian: PSDMatrix, /) -> Tensor:
         # Move all computations on cpu to avoid moving memory between cpu and gpu at each iteration
         device = gramian.device
         dtype = gramian.dtype
diff --git a/src/torchjd/aggregation/_random.py b/src/torchjd/aggregation/_random.py
@@ -26,7 +26,7 @@ class RandomWeighting(Weighting[Matrix]):
     at each call.
     """
 
-    def forward(self, matrix: Tensor) -> Tensor:
+    def forward(self, matrix: Tensor, /) -> Tensor:
         random_vector = torch.randn(matrix.shape[0], device=matrix.device, dtype=matrix.dtype)
         weights = F.softmax(random_vector, dim=-1)
         return weights
diff --git a/src/torchjd/aggregation/_sum.py b/src/torchjd/aggregation/_sum.py
@@ -23,7 +23,7 @@ class SumWeighting(Weighting[Matrix]):
     :math:`\begin{bmatrix} 1 & \dots & 1 \end{bmatrix}^T \in \mathbb{R}^m`.
     """
 
-    def forward(self, matrix: Tensor) -> Tensor:
+    def forward(self, matrix: Tensor, /) -> Tensor:
         device = matrix.device
         dtype = matrix.dtype
         weights = torch.ones(matrix.shape[0], device=device, dtype=dtype)
diff --git a/src/torchjd/aggregation/_upgrad.py b/src/torchjd/aggregation/_upgrad.py
@@ -86,7 +86,7 @@ def __init__(
         self.reg_eps = reg_eps
         self.solver: SUPPORTED_SOLVER = solver
 
-    def forward(self, gramian: PSDMatrix) -> Tensor:
+    def forward(self, gramian: PSDMatrix, /) -> Tensor:
         U = torch.diag(self.weighting(gramian))
         G = regularize(normalize(gramian, self.norm_eps), self.reg_eps)
         W = project_weights(U, G, self.solver)
diff --git a/src/torchjd/aggregation/_weighting_bases.py b/src/torchjd/aggregation/_weighting_bases.py
@@ -24,7 +24,7 @@ def __init__(self):
         super().__init__()
 
     @abstractmethod
-    def forward(self, stat: _T) -> Tensor:
+    def forward(self, stat: _T, /) -> Tensor:
         """Computes the vector of weights from the input stat."""
 
     def __call__(self, stat: Tensor) -> Tensor:
@@ -51,7 +51,7 @@ def __init__(self, weighting: Weighting[_FnOutputT], fn: Callable[[_T], _FnOutpu
         self.fn = fn
         self.weighting = weighting
 
-    def forward(self, stat: _T) -> Tensor:
+    def forward(self, stat: _T, /) -> Tensor:
         return self.weighting(self.fn(stat))
 
 
diff --git a/src/torchjd/autogram/_jacobian_computer.py b/src/torchjd/autogram/_jacobian_computer.py
@@ -56,6 +56,7 @@ def _compute_jacobian(
         grad_outputs: tuple[Tensor, ...],
         args: tuple[PyTree, ...],
         kwargs: dict[str, PyTree],
+        /,
     ) -> Matrix:
         """
         Computes and returns the Jacobian. The output must be a matrix (2D Tensor).
@@ -75,6 +76,7 @@ def _compute_jacobian(
         grad_outputs: tuple[Tensor, ...],
         args: tuple[PyTree, ...],
         kwargs: dict[str, PyTree],
+        /,
     ) -> Matrix:
         grad_outputs_in_dims = (0,) * len(grad_outputs)
         args_in_dims = tree_map(lambda t: 0 if isinstance(t, Tensor) else None, args)
@@ -133,6 +135,7 @@ def _compute_jacobian(
         grad_outputs: tuple[Tensor, ...],
         _: tuple[PyTree, ...],
         __: dict[str, PyTree],
+        /,
     ) -> Matrix:
         flat_rg_params, ___ = tree_flatten(self.rg_params)
         grads = torch.autograd.grad(
@@ -172,7 +175,7 @@ def vmap(
         jac_outputs: tuple[Tensor, ...],
         args: tuple[PyTree, ...],
         kwargs: dict[str, PyTree],
-    ) -> tuple[Tensor, None]:
+    ) -> tuple[Tensor, None]:  # type: ignore[reportIncompatibleMethodOverride]
         # There is a non-batched dimension
         # We do not vmap over the args, kwargs, or rg_outputs for the non-batched dimension
         generalized_jacobian = torch.vmap(compute_jacobian_fn, in_dims=in_dims[1:])(
diff --git a/src/torchjd/autogram/_module_hook_manager.py b/src/torchjd/autogram/_module_hook_manager.py
@@ -173,7 +173,7 @@ def setup_context(
         ctx,
         inputs: tuple,
         _,
-    ) -> None:
+    ) -> None:  # type: ignore[reportIncompatibleMethodOverride]
         ctx.gramian_accumulation_phase = inputs[0]
         ctx.gramian_computer = inputs[1]
         ctx.args = inputs[2]
diff --git a/src/torchjd/autojac/_transform/_accumulate.py b/src/torchjd/autojac/_transform/_accumulate.py
@@ -13,7 +13,7 @@ class AccumulateGrad(Transform):
     should not be used elsewhere.
     """
 
-    def __call__(self, gradients: TensorDict) -> TensorDict:
+    def __call__(self, gradients: TensorDict, /) -> TensorDict:
         accumulate_grads(gradients.keys(), gradients.values())
         return {}
 
@@ -30,7 +30,7 @@ class AccumulateJac(Transform):
     should not be used elsewhere.
     """
 
-    def __call__(self, jacobians: TensorDict) -> TensorDict:
+    def __call__(self, jacobians: TensorDict, /) -> TensorDict:
         accumulate_jacs(jacobians.keys(), jacobians.values())
         return {}
 
diff --git a/src/torchjd/autojac/_transform/_base.py b/src/torchjd/autojac/_transform/_base.py
@@ -41,7 +41,7 @@ def __str__(self) -> str:
         return type(self).__name__
 
     @abstractmethod
-    def __call__(self, input: TensorDict) -> TensorDict:
+    def __call__(self, input: TensorDict, /) -> TensorDict:
         """Applies the transform to the input."""
 
     @abstractmethod
@@ -76,7 +76,7 @@ def __init__(self, outer: Transform, inner: Transform):
     def __str__(self) -> str:
         return str(self.outer) + " ∘ " + str(self.inner)
 
-    def __call__(self, input: TensorDict) -> TensorDict:
+    def __call__(self, input: TensorDict, /) -> TensorDict:
         intermediate = self.inner(input)
         return self.outer(intermediate)
 
@@ -107,7 +107,7 @@ def __str__(self) -> str:
                 strings.append(s)
         return "(" + " | ".join(strings) + ")"
 
-    def __call__(self, tensor_dict: TensorDict) -> TensorDict:
+    def __call__(self, tensor_dict: TensorDict, /) -> TensorDict:
         union: TensorDict = {}
         for transform in self.transforms:
             union |= transform(tensor_dict)
diff --git a/src/torchjd/autojac/_transform/_diagonalize.py b/src/torchjd/autojac/_transform/_diagonalize.py
@@ -60,7 +60,7 @@ def __init__(self, key_order: OrderedSet[Tensor]):
             self.indices.append((begin, end))
             begin = end
 
-    def __call__(self, tensors: TensorDict) -> TensorDict:
+    def __call__(self, tensors: TensorDict, /) -> TensorDict:
         flattened_considered_values = [tensors[key].reshape([-1]) for key in self.key_order]
         diagonal_matrix = torch.cat(flattened_considered_values).diag()
         diagonalized_tensors = {
diff --git a/src/torchjd/autojac/_transform/_differentiate.py b/src/torchjd/autojac/_transform/_differentiate.py
@@ -37,15 +37,15 @@ def __init__(
         self.retain_graph = retain_graph
         self.create_graph = create_graph
 
-    def __call__(self, tensors: TensorDict) -> TensorDict:
+    def __call__(self, tensors: TensorDict, /) -> TensorDict:
         tensor_outputs = [tensors[output] for output in self.outputs]
 
         differentiated_tuple = self._differentiate(tensor_outputs)
         new_differentiations = dict(zip(self.inputs, differentiated_tuple))
         return type(tensors)(new_differentiations)
 
     @abstractmethod
-    def _differentiate(self, tensor_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]:
+    def _differentiate(self, tensor_outputs: Sequence[Tensor], /) -> tuple[Tensor, ...]:
         """
         Abstract method for differentiating the outputs with respect to the inputs, and applying the
         linear transformations represented by the tensor_outputs to the results.
diff --git a/src/torchjd/autojac/_transform/_grad.py b/src/torchjd/autojac/_transform/_grad.py
@@ -34,7 +34,7 @@ def __init__(
     ):
         super().__init__(outputs, inputs, retain_graph, create_graph)
 
-    def _differentiate(self, grad_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]:
+    def _differentiate(self, grad_outputs: Sequence[Tensor], /) -> tuple[Tensor, ...]:
         """
         Computes the gradient of each output element with respect to each input tensor, and applies
         the linear transformations represented by the grad_outputs to the results.
diff --git a/src/torchjd/autojac/_transform/_init.py b/src/torchjd/autojac/_transform/_init.py
@@ -16,7 +16,7 @@ class Init(Transform):
     def __init__(self, values: Set[Tensor]):
         self.values = values
 
-    def __call__(self, input: TensorDict) -> TensorDict:
+    def __call__(self, input: TensorDict, /) -> TensorDict:
         return {value: torch.ones_like(value) for value in self.values}
 
     def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
diff --git a/src/torchjd/autojac/_transform/_jac.py b/src/torchjd/autojac/_transform/_jac.py
@@ -42,7 +42,7 @@ def __init__(
         super().__init__(outputs, inputs, retain_graph, create_graph)
         self.chunk_size = chunk_size
 
-    def _differentiate(self, jac_outputs: Sequence[Tensor]) -> tuple[Tensor, ...]:
+    def _differentiate(self, jac_outputs: Sequence[Tensor], /) -> tuple[Tensor, ...]:
         """
         Computes the jacobian of each output with respect to each input, and applies the linear
         transformations represented by the jac_outputs to the results.
diff --git a/src/torchjd/autojac/_transform/_ordered_set.py b/src/torchjd/autojac/_transform/_ordered_set.py
@@ -20,10 +20,10 @@ def difference_update(self, elements: set[_T]) -> None:
         for element in elements:
             self.discard(element)
 
-    def add(self, element: _T) -> None:
+    def add(self, value: _T) -> None:
         """Adds the specified element to the OrderedSet."""
 
-        self.ordered_dict[element] = None
+        self.ordered_dict[value] = None
 
     def __add__(self, other: OrderedSet[_T]) -> OrderedSet[_T]:
         """Creates a new OrderedSet with the elements of self followed by the elements of other."""
@@ -40,5 +40,5 @@ def __iter__(self) -> Iterator[_T]:
     def __len__(self) -> int:
         return len(self.ordered_dict)
 
-    def __contains__(self, element: object) -> bool:
-        return element in self.ordered_dict
+    def __contains__(self, x: object) -> bool:
+        return x in self.ordered_dict
diff --git a/src/torchjd/autojac/_transform/_select.py b/src/torchjd/autojac/_transform/_select.py
@@ -15,7 +15,7 @@ class Select(Transform):
     def __init__(self, keys: Set[Tensor]):
         self.keys = keys
 
-    def __call__(self, tensor_dict: TensorDict) -> TensorDict:
+    def __call__(self, tensor_dict: TensorDict, /) -> TensorDict:
         output = {key: tensor_dict[key] for key in self.keys}
         return type(tensor_dict)(output)
 
diff --git a/src/torchjd/autojac/_transform/_stack.py b/src/torchjd/autojac/_transform/_stack.py
diff --git a/tests/doc/test_rst.py b/tests/doc/test_rst.py
diff --git a/tests/plots/interactive_plotter.py b/tests/plots/interactive_plotter.py
diff --git a/tests/unit/autogram/test_engine.py b/tests/unit/autogram/test_engine.py
diff --git a/tests/unit/autojac/_transform/test_base.py b/tests/unit/autojac/_transform/test_base.py
diff --git a/tests/unit/autojac/_transform/test_stack.py b/tests/unit/autojac/_transform/test_stack.py
diff --git a/tests/unit/autojac/test_backward.py b/tests/unit/autojac/test_backward.py
diff --git a/tests/unit/autojac/test_jac.py b/tests/unit/autojac/test_jac.py
diff --git a/tests/unit/autojac/test_mtl_backward.py b/tests/unit/autojac/test_mtl_backward.py
diff --git a/tests/utils/architectures.py b/tests/utils/architectures.py

Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ Changelog = "https://github.com/TorchJD/torchjd/blob/main/CHANGELOG.md"`
`66`	`66`
`67`	`67`	`[dependency-groups]`
`68`	`68`	`check = [`
`69`		`- "mypy>=1.16.0",`
	`69`	`+ "ty>=0.0.14",`
`70`	`70`	`"pre-commit>=2.9.2", # isort doesn't work before 2.9.2`
`71`	`71`	`]`
`72`	`72`
`@@ -114,3 +114,7 @@ exclude_lines = [`
`114`	`114`	`"pragma: not covered",`
`115`	`115`	`"@overload",`
`116`	`116`	`]`
	`117`	`+`
	`118`	`+[tool.ty.src]`
	`119`	`+include = ["src", "tests"]`
	`120`	`+exclude = ["src/torchjd/aggregation/_nash_mtl.py"]`