refactor(autojac): Reorder create_jac_dict parameters

Author: PierreQuinton

src/torchjd/autojac/_backward.py

@@ -116,7 +116,7 @@ def backward(
     else:
         inputs_ = OrderedSet(inputs)
 
-    jac_tensors_dict = create_jac_dict(tensors_, jac_tensors, "jac_tensors", "tensors")
+    jac_tensors_dict = create_jac_dict(tensors_, jac_tensors, "tensors", "jac_tensors")
     transform = _create_transform(tensors_, inputs_, parallel_chunk_size, retain_graph)
     transform(jac_tensors_dict)
 

src/torchjd/autojac/_jac.py

@@ -155,7 +155,7 @@ def jac(
     inputs_with_repetition = cast(Sequence[Tensor], (inputs,) if is_tensor_like(inputs) else inputs)
     inputs_ = OrderedSet(inputs_with_repetition)
 
-    jac_outputs_dict = create_jac_dict(outputs_, jac_outputs, "jac_outputs", "outputs")
+    jac_outputs_dict = create_jac_dict(outputs_, jac_outputs, "outputs", "jac_outputs")
     transform = _create_transform(outputs_, inputs_, parallel_chunk_size, retain_graph)
     result = transform(jac_outputs_dict)
     return tuple(result[input] for input in inputs_with_repetition)

src/torchjd/autojac/_utils.py

@@ -106,8 +106,8 @@ def check_matching_grad_shapes(
 def create_jac_dict(
     tensors: OrderedSet[Tensor],
     opt_jacobians: Sequence[Tensor] | Tensor | None,
-    jacobian_param_name: str,
     tensor_param_name: str,
+    jacobian_param_name: str,
 ) -> dict[Tensor, Tensor]:
     """
     Creates a dictionary mapping tensors to their corresponding Jacobians.
@@ -118,8 +118,8 @@ def create_jac_dict(
     :param tensors: The tensors to differentiate.
     :param opt_jacobians: The initial Jacobians to backpropagate. If ``None``, defaults to
         identity.
-    :param jacobian_param_name: The name of the jacobian parameter for error messages.
     :param tensor_param_name: The name of the tensor parameter for error messages.
+    :param jacobian_param_name: The name of the jacobian parameter for error messages.
     """
     from torchjd.autojac._transform._diagonalize import Diagonalize
     from torchjd.autojac._transform._init import Init

tests/unit/autojac/test_backward.py

@@ -26,8 +26,8 @@ def test_check_create_transform(default_jac_tensors: bool) -> None:
     jac_tensors = create_jac_dict(
         tensors=OrderedSet([y1, y2]),
         opt_jacobians=optional_jac_tensors,
-        jacobian_param_name="jac_tensors",
         tensor_param_name="tensors",
+        jacobian_param_name="jac_tensors",
     )
     transform = _create_transform(
         tensors=OrderedSet([y1, y2]),

tests/unit/autojac/test_jac.py

@@ -26,8 +26,8 @@ def test_check_create_transform(default_jac_outputs: bool) -> None:
     jac_outputs = create_jac_dict(
         tensors=OrderedSet([y1, y2]),
         opt_jacobians=optional_jac_outputs,
-        jacobian_param_name="jac_outputs",
         tensor_param_name="outputs",
+        jacobian_param_name="jac_outputs",
     )
     transform = _create_transform(
         outputs=OrderedSet([y1, y2]),