style(autojac): Add missing type hints (#367)

* Add type hint to union in _stack
* Add type hint to result in _get_descendant_accumulate_grads
* Be more specific about the type of Select when instantiating it

Author: ValerianRey

src/torchjd/_autojac/_mtl_backward.py

@@ -179,10 +179,10 @@ def _create_task_transform(
 
     # Transform that accumulates the gradients w.r.t. the task-specific parameters into their
     # .grad fields.
-    accumulate = Accumulate() << Select(task_params)
+    accumulate = Accumulate() << Select[Gradients](task_params)
 
     # Transform that backpropagates the gradients of the losses w.r.t. the features.
-    backpropagate = Select(features)
+    backpropagate = Select[Gradients](features)
 
     # Transform that accumulates the gradient of the losses w.r.t. the task-specific parameters into
     # their .grad fields and backpropagates the gradient of the losses w.r.t. to the features.

src/torchjd/_autojac/_transform/_stack.py

@@ -36,7 +36,7 @@ def _stack(gradient_dicts: list[Gradients]) -> Jacobians:
     # It is important to first remove duplicate keys before computing their associated
     # stacked tensor. Otherwise, some computations would be duplicated. Therefore, we first compute
     # unique_keys, and only then, we compute the stacked tensors.
-    union = {}
+    union: dict[Tensor, Tensor] = {}
     for d in gradient_dicts:
         union |= d
     unique_keys = union.keys()

src/torchjd/_autojac/_utils.py

@@ -67,7 +67,7 @@ def _get_descendant_accumulate_grads(
     """
 
     excluded_nodes = set(excluded_nodes)  # Re-instantiate set to avoid modifying input
-    result = OrderedSet([])
+    result: OrderedSet[Node] = OrderedSet([])
     roots.difference_update(excluded_nodes)
     nodes_to_traverse = deque(roots)