refactor(autojac): Fix genericity of Transforms (#369)

* Make _B contravariant and _C covariant
* Fix generic parametrization in Composition and Conjunction
* Make Stack contravariant in its input type

Author: ValerianRey

src/torchjd/_autojac/_transform/_base.py

@@ -51,22 +51,22 @@ def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
     __or__ = conjunct
 
 
-class Composition(Transform[_A, _C]):
+class Composition(Transform[_B, _C]):
     """
     Transform corresponding to the composition of two transforms inner and outer.
 
     :param inner: The transform to apply first, to the input.
     :param outer: The transform to apply second, to the result of ``inner``.
     """
 
-    def __init__(self, outer: Transform[_B, _C], inner: Transform[_A, _B]):
+    def __init__(self, outer: Transform[_A, _C], inner: Transform[_B, _A]):
         self.outer = outer
         self.inner = inner
 
     def __str__(self) -> str:
         return str(self.outer) + " ∘ " + str(self.inner)
 
-    def __call__(self, input: _A) -> _C:
+    def __call__(self, input: _B) -> _C:
         intermediate = self.inner(input)
         return self.outer(intermediate)
 
@@ -76,15 +76,15 @@ def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
         return output_keys
 
 
-class Conjunction(Transform[_A, _B]):
+class Conjunction(Transform[_B, _C]):
     """
     Transform applying several transforms to the same input, and combining the results (by union)
     into a single TensorDict.
 
     :param transforms: The transforms to apply. Their outputs should have disjoint sets of keys.
     """
 
-    def __init__(self, transforms: Sequence[Transform[_A, _B]]):
+    def __init__(self, transforms: Sequence[Transform[_B, _C]]):
         self.transforms = transforms
 
     def __str__(self) -> str:
@@ -97,10 +97,10 @@ def __str__(self) -> str:
                 strings.append(s)
         return "(" + " | ".join(strings) + ")"
 
-    def __call__(self, tensor_dict: _A) -> _B:
+    def __call__(self, tensor_dict: _B) -> _C:
         tensor_dicts = [transform(tensor_dict) for transform in self.transforms]
-        output_type: type[_A] = EmptyTensorDict
-        output: _A = EmptyTensorDict()
+        output_type: type[_B] = EmptyTensorDict
+        output: _B = EmptyTensorDict()
         for tensor_dict in tensor_dicts:
             output_type = _least_common_ancestor(output_type, type(tensor_dict))
             output |= tensor_dict

src/torchjd/_autojac/_transform/_stack.py

@@ -5,10 +5,10 @@
 
 from ._base import Transform
 from ._materialize import materialize
-from ._tensor_dict import _A, Gradients, Jacobians
+from ._tensor_dict import _B, Gradients, Jacobians
 
 
-class Stack(Transform[_A, Jacobians]):
+class Stack(Transform[_B, Jacobians]):
     """
     Transform applying several transforms to the same input, and combining the results (by stacking)
     into a single TensorDict.
@@ -20,10 +20,10 @@ class Stack(Transform[_A, Jacobians]):
         at the positions corresponding to those dicts.
     """
 
-    def __init__(self, transforms: Sequence[Transform[_A, Gradients]]):
+    def __init__(self, transforms: Sequence[Transform[_B, Gradients]]):
         self.transforms = transforms
 
-    def __call__(self, input: _A) -> Jacobians:
+    def __call__(self, input: _B) -> Jacobians:
         results = [transform(input) for transform in self.transforms]
         result = _stack(results)
         return result

src/torchjd/_autojac/_transform/_tensor_dict.py

@@ -182,5 +182,5 @@ def _check_corresponding_numel(key: Tensor, value: Tensor, dim: int) -> None:
 
 
 _A = TypeVar("_A", bound=TensorDict)
-_B = TypeVar("_B", bound=TensorDict)
-_C = TypeVar("_C", bound=TensorDict)
+_B = TypeVar("_B", bound=TensorDict, contravariant=True)
+_C = TypeVar("_C", bound=TensorDict, covariant=True)