Make Transform subclasses also have pos-only arguments

Similarly to the previous commit, this is not strictly necessary according to LSP, but I think it's weird that subclasses don't enforce pos-only arguments if the parent class enforces that.

Author: ValerianRey

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

@@ -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,7 +37,7 @@ 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)

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

src/torchjd/autojac/_transform/_stack.py

@@ -23,7 +23,7 @@ class Stack(Transform):
     def __init__(self, transforms: Sequence[Transform]):
         self.transforms = transforms
 
-    def __call__(self, input: TensorDict) -> TensorDict:
+    def __call__(self, input: TensorDict, /) -> TensorDict:
         results = [transform(input) for transform in self.transforms]
         result = _stack(results)
         return result

tests/unit/autojac/_transform/test_base.py

@@ -17,7 +17,7 @@ def __init__(self, required_keys: set[Tensor], output_keys: set[Tensor]):
     def __str__(self):
         return "T"
 
-    def __call__(self, input: TensorDict) -> TensorDict:
+    def __call__(self, input: TensorDict, /) -> TensorDict:
         # Ignore the input, create a dictionary with the right keys as an output.
         output_dict = {key: empty_(0) for key in self._output_keys}
         return output_dict

tests/unit/autojac/_transform/test_stack.py

@@ -15,7 +15,7 @@ class FakeGradientsTransform(Transform):
     def __init__(self, keys: Iterable[Tensor]):
         self.keys = set(keys)
 
-    def __call__(self, input: TensorDict) -> TensorDict:
+    def __call__(self, input: TensorDict, /) -> TensorDict:
         return {key: torch.ones_like(key) for key in self.keys}
 
     def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]: