From 5368378850333ad93db4b62c523568693ed7c5b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Val=C3=A9rian=20Rey?= Date: Wed, 7 Oct 2026 17:49:02 +0200 Subject: [PATCH] Simplify autojac input utilities `as_checked_ordered_set` built the same OrderedSet twice: once to check uniqueness, and once to return it. `create_jac_dict` imported Diagonalize and Init inside the function, although the module already imports `._transform` at the top, so they can be imported there. Co-Authored-By: Claude Opus 5.5 --- src/torchjd/autojac/_utils.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/torchjd/autojac/_utils.py b/src/torchjd/autojac/_utils.py index 56d67ff8a..4946b5de9 100644 --- a/src/torchjd/autojac/_utils.py +++ b/src/torchjd/autojac/_utils.py @@ -6,7 +6,7 @@ from torch.autograd.graph import Node from torch.overrides import is_tensor_like -from ._transform import OrderedSet +from ._transform import Diagonalize, Init, OrderedSet def check_optional_positive_chunk_size(parallel_chunk_size: int | None) -> None: @@ -30,7 +30,7 @@ def as_checked_ordered_set( if len(output) != original_length: raise ValueError(f"`{variable_name}` should contain unique elements.") - return OrderedSet(tensors) + return output def check_matching_length( @@ -121,9 +121,6 @@ def create_jac_dict( :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 - if opt_jacobians is None: init = Init(tensors) diag = Diagonalize(tensors)