typing: Replace mypy with ty

.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

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

@@ -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"]

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

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
 

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)

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()

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 = (

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]

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

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()
 

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

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

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)

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)

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))
 
 

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:])(

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]

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 {}
 

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)

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 = {

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.

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.

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]:

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.

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

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)