refactor: Add casts when necessary (#372)

* Add cast in get_leaf_tensors
* Add cast of w_opt to np.ndarray in _CAGradWeighting.forward
* Add cast to type[TensorDict] in _least_common_ancestor

Author: ValerianRey

src/torchjd/_autojac/_transform/_tensor_dict.py

@@ -1,4 +1,4 @@
-from typing import TypeVar
+from typing import TypeVar, cast
 
 from torch import Tensor
 
@@ -131,7 +131,7 @@ def _least_common_ancestor(first: type[TensorDict], second: type[TensorDict]) ->
     output = TensorDict
     for candidate_type in first_mro:
         if issubclass(second, candidate_type):
-            output = candidate_type
+            output = cast(type[TensorDict], candidate_type)
             break
     return output
 

src/torchjd/_autojac/_utils.py

@@ -1,5 +1,6 @@
 from collections import deque
 from collections.abc import Iterable, Sequence
+from typing import cast
 
 from torch import Tensor
 from torch.autograd.graph import Node
@@ -48,8 +49,8 @@ def get_leaf_tensors(tensors: Iterable[Tensor], excluded: Iterable[Tensor]) -> O
         raise ValueError("All `excluded` tensors should have a `grad_fn`.")
 
     accumulate_grads = _get_descendant_accumulate_grads(
-        roots=OrderedSet([tensor.grad_fn for tensor in tensors]),
-        excluded_nodes={tensor.grad_fn for tensor in excluded},
+        roots=cast(OrderedSet[Node], OrderedSet([tensor.grad_fn for tensor in tensors])),
+        excluded_nodes=cast(set[Node], {tensor.grad_fn for tensor in excluded}),
     )
 
     # accumulate_grads contains instances of AccumulateGrad, which contain a `variable` field.

src/torchjd/aggregation/_cagrad.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 from ._utils.check_dependencies import check_dependencies_are_installed
 from ._weighting_bases import PSDMatrix, Weighting
 
@@ -101,7 +103,7 @@ def forward(self, gramian: Tensor) -> Tensor:
         problem = cp.Problem(objective=cp.Minimize(cost), constraints=[w >= 0, cp.sum(w) == 1])
 
         problem.solve(cp.CLARABEL)
-        w_opt = w.value
+        w_opt = cast(np.ndarray, w.value)
 
         g_w_norm = np.linalg.norm(reduced_array.T @ w_opt, 2).item()
         if g_w_norm >= self.norm_eps: