Add FBT (boolean traps)

PierreQuinton · PierreQuinton · commit ef0516f6e5ce · 2026-02-05T13:40:16.000+01:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -134,6 +134,7 @@ select = [
     "W",    # pycodestyle Warning
     "I",    # isort
     "UP",   # pyupgrade
+    "FBT",  # flake8-boolean-trap
     "B",    # flake8-bugbear
     "C4",   # flake8-comprehensions
     "FIX",  # flake8-fixme
diff --git a/src/torchjd/autogram/_engine.py b/src/torchjd/autogram/_engine.py
@@ -306,7 +306,7 @@ def compute_gramian(self, output: Tensor) -> Tensor:
 
         return gramian
 
-    def _compute_square_gramian(self, output: Tensor, has_non_batch_dim: bool) -> PSDMatrix:
+    def _compute_square_gramian(self, output: Tensor, *, has_non_batch_dim: bool) -> PSDMatrix:
         leaf_targets = list(self._target_edges.get_leaf_edges({get_gradient_edge(output)}))
 
         def differentiation(_grad_output: Tensor) -> tuple[Tensor, ...]:
diff --git a/src/torchjd/autogram/_module_hook_manager.py b/src/torchjd/autogram/_module_hook_manager.py
@@ -36,7 +36,7 @@ def __init__(
     ):
         self._target_edges = target_edges
         self._gramian_accumulator = gramian_accumulator
-        self.gramian_accumulation_phase = BoolRef(False)
+        self.gramian_accumulation_phase = BoolRef(value=False)
         self._handles: list[TorchRemovableHandle] = []
 
         # When the ModuleHookManager is not referenced anymore, there is no reason to keep the hooks
@@ -79,7 +79,7 @@ def remove_hooks(handles: list[TorchRemovableHandle]) -> None:
 class BoolRef:
     """Class wrapping a boolean value, acting as a reference to this boolean value."""
 
-    def __init__(self, value: bool):
+    def __init__(self, *, value: bool):
         self.value = value
 
     def __bool__(self) -> bool:
diff --git a/src/torchjd/autojac/_backward.py b/src/torchjd/autojac/_backward.py
@@ -9,6 +9,7 @@
 def backward(
     tensors: Sequence[Tensor] | Tensor,
     inputs: Iterable[Tensor] | None = None,
+    *,
     retain_graph: bool = False,
     parallel_chunk_size: int | None = None,
 ) -> None:
@@ -86,6 +87,7 @@ def backward(
 def _create_transform(
     tensors: OrderedSet[Tensor],
     inputs: OrderedSet[Tensor],
+    *,
     retain_graph: bool,
     parallel_chunk_size: int | None,
 ) -> Transform:
diff --git a/src/torchjd/autojac/_jac.py b/src/torchjd/autojac/_jac.py
@@ -17,6 +17,7 @@
 def jac(
     outputs: Sequence[Tensor] | Tensor,
     inputs: Iterable[Tensor] | None = None,
+    *,
     retain_graph: bool = False,
     parallel_chunk_size: int | None = None,
 ) -> tuple[Tensor, ...]:
@@ -136,6 +137,7 @@ def jac(
 def _create_transform(
     outputs: OrderedSet[Tensor],
     inputs: OrderedSet[Tensor],
+    *,
     retain_graph: bool,
     parallel_chunk_size: int | None,
 ) -> Transform:
diff --git a/src/torchjd/autojac/_jac_to_grad.py b/src/torchjd/autojac/_jac_to_grad.py
@@ -11,6 +11,7 @@
 def jac_to_grad(
     tensors: Iterable[Tensor],
     aggregator: Aggregator,
+    *,
     retain_jac: bool = False,
 ) -> None:
     r"""
diff --git a/src/torchjd/autojac/_mtl_backward.py b/src/torchjd/autojac/_mtl_backward.py
@@ -21,6 +21,7 @@ def mtl_backward(
     features: Sequence[Tensor] | Tensor,
     tasks_params: Sequence[Iterable[Tensor]] | None = None,
     shared_params: Iterable[Tensor] | None = None,
+    *,
     retain_graph: bool = False,
     parallel_chunk_size: int | None = None,
 ) -> None:
@@ -113,6 +114,7 @@ def _create_transform(
     features: OrderedSet[Tensor],
     tasks_params: list[OrderedSet[Tensor]],
     shared_params: OrderedSet[Tensor],
+    *,
     retain_graph: bool,
     parallel_chunk_size: int | None,
 ) -> Transform:
@@ -152,6 +154,7 @@ def _create_task_transform(
     features: OrderedSet[Tensor],
     task_params: OrderedSet[Tensor],
     loss: OrderedSet[Tensor],  # contains a single scalar loss
+    *,
     retain_graph: bool,
 ) -> Transform:
     # Tensors with respect to which we compute the gradients.
diff --git a/src/torchjd/autojac/_transform/_differentiate.py b/src/torchjd/autojac/_transform/_differentiate.py
@@ -29,6 +29,7 @@ def __init__(
         self,
         outputs: OrderedSet[Tensor],
         inputs: OrderedSet[Tensor],
+        *,
         retain_graph: bool,
         create_graph: bool,
     ):
@@ -64,7 +65,7 @@ def check_keys(self, input_keys: set[Tensor], /) -> set[Tensor]:
             )
         return set(self.inputs)
 
-    def _get_vjp(self, grad_outputs: Sequence[Tensor], retain_graph: bool) -> tuple[Tensor, ...]:
+    def _get_vjp(self, grad_outputs: Sequence[Tensor], *, retain_graph: bool) -> tuple[Tensor, ...]:
         optional_grads = torch.autograd.grad(
             self.outputs,
             self.inputs,
diff --git a/src/torchjd/autojac/_transform/_grad.py b/src/torchjd/autojac/_transform/_grad.py
@@ -29,6 +29,7 @@ def __init__(
         self,
         outputs: OrderedSet[Tensor],
         inputs: OrderedSet[Tensor],
+        *,
         retain_graph: bool = False,
         create_graph: bool = False,
     ):
diff --git a/src/torchjd/autojac/_transform/_jac.py b/src/torchjd/autojac/_transform/_jac.py
@@ -35,6 +35,7 @@ def __init__(
         self,
         outputs: OrderedSet[Tensor],
         inputs: OrderedSet[Tensor],
+        *,
         chunk_size: int | None,
         retain_graph: bool = False,
         create_graph: bool = False,
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -27,7 +27,7 @@ def fix_randomness() -> None:
     # reproducibility on GPU. We also use GPU to benchmark algorithms, and we would rather have them
     # use non-deterministic but faster algorithms.
     if DEVICE.type == "cpu":
-        torch.use_deterministic_algorithms(True)
+        torch.use_deterministic_algorithms(mode=True)
 
 
 def pytest_addoption(parser):
diff --git a/tests/plots/_utils.py b/tests/plots/_utils.py
@@ -59,6 +59,7 @@ def make_vector_scatter(
     gradient: torch.Tensor,
     color: str,
     label: str,
+    *,
     showlegend: bool = False,
     dash: bool = False,
     textposition: str = "bottom center",
@@ -94,6 +95,7 @@ def make_cone_scatter(
     opening: float,
     label: str,
     scale: float = 100.0,
+    *,
     printable: bool = False,
 ) -> Scatter:
     if opening < -1e-8:
@@ -171,6 +173,7 @@ def make_polygon_scatter(points: list[torch.Tensor]) -> Scatter:
 def make_right_angle(
     vector: torch.Tensor,
     size: float,
+    *,
     positive_para: bool = True,
     positive_orth: bool = True,
 ) -> list[torch.Tensor]:
diff --git a/tests/profiling/plot_memory_timeline.py b/tests/profiling/plot_memory_timeline.py
@@ -62,7 +62,7 @@ def plot_memory_timelines(experiment: str, folders: list[str]) -> None:
     ax.set_ylabel("Total Allocated (bytes)", fontsize=12)
     ax.set_title(f"Memory Timeline: {experiment}", fontsize=14, fontweight="bold")
     ax.legend(loc="best", fontsize=11)
-    ax.grid(True, alpha=0.3)
+    ax.grid(visible=True, alpha=0.3)
     ax.set_ylim(bottom=0)
     fig.tight_layout()
 
diff --git a/tests/unit/autogram/test_engine.py b/tests/unit/autogram/test_engine.py
@@ -361,6 +361,7 @@ def test_iwrm_steps_with_autogram(factory: ModuleFactory, batch_size: int, batch
 def test_autograd_while_modules_are_hooked(
     factory: ModuleFactory,
     batch_size: int,
+    *,
     use_engine: bool,
     batch_dim: int | None,
 ):
diff --git a/tests/unit/autojac/test_backward.py b/tests/unit/autojac/test_backward.py
@@ -48,7 +48,7 @@ def test_jac_is_populated():
 @mark.parametrize("chunk_size", [1, 2, None])
 def test_value_is_correct(
     shape: tuple[int, int],
-    manually_specify_inputs: bool,
+    *manually_specify_inputs: bool,
     chunk_size: int | None,
 ):
     """
diff --git a/tests/unit/autojac/test_jac.py b/tests/unit/autojac/test_jac.py
@@ -52,6 +52,7 @@ def test_jac():
 @mark.parametrize("chunk_size", [1, 2, None])
 def test_value_is_correct(
     shape: tuple[int, int],
+    *,
     manually_specify_inputs: bool,
     chunk_size: int | None,
 ):
diff --git a/tests/unit/autojac/test_jac_to_grad.py b/tests/unit/autojac/test_jac_to_grad.py
@@ -86,7 +86,7 @@ def test_no_tensors():
 
 
 @mark.parametrize("retain_jac", [True, False])
-def test_jacs_are_freed(retain_jac: bool):
+def test_jacs_are_freed(*, retain_jac: bool):
     """Tests that jac_to_grad frees the jac fields if an only if retain_jac is False."""
 
     aggregator = UPGrad()
diff --git a/tests/unit/autojac/test_mtl_backward.py b/tests/unit/autojac/test_mtl_backward.py
@@ -67,6 +67,7 @@ def test_shape_is_correct():
 @mark.parametrize("chunk_size", [1, 2, None])
 def test_value_is_correct(
     shape: tuple[int, int],
+    *,
     manually_specify_shared_params: bool,
     manually_specify_tasks_params: bool,
     chunk_size: int | None,
diff --git a/tests/utils/architectures.py b/tests/utils/architectures.py
@@ -524,7 +524,7 @@ def __init__(self):
         super().__init__()
         self.non_frozen = nn.Linear(50, 10)
         self.all_frozen = nn.Linear(50, 10)
-        self.all_frozen.requires_grad_(False)
+        self.all_frozen.requires_grad_(requires_grad=False)
 
     def forward(self, input: Tensor) -> Tensor:
         return self.all_frozen(input) + self.non_frozen(input**2 / 5.0)
diff --git a/tests/utils/forward_backwards.py b/tests/utils/forward_backwards.py
@@ -114,6 +114,7 @@ def reshape_raw_losses(raw_losses: Tensor) -> Tensor:
 def compute_gramian_with_autograd(
     output: Tensor,
     params: list[nn.Parameter],
+    *,
     retain_graph: bool = False,
 ) -> PSDTensor:
     """
