Fix docstring formatting

Author: ValerianRey

tests/unit/autojac/_transform/test_base.py

@@ -35,8 +35,8 @@ def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
 
 def test_composition_check_keys():
     """
-    Tests that `check_keys` works correctly for a composition of transforms: the inner
-    transform's `output_keys` has to match with the outer transform's `required_keys`.
+    Tests that `check_keys` works correctly for a composition of transforms: the inner transform's
+    `output_keys` has to match with the outer transform's `required_keys`.
     """
 
     a1 = torch.randn([2])
@@ -54,8 +54,8 @@ def test_composition_check_keys():
 
 def test_conjunct_check_keys_1():
     """
-    Tests that `check_keys` works correctly for a conjunction of transforms: all transforms
-    should successfully check their keys.
+    Tests that `check_keys` works correctly for a conjunction of transforms: all transforms should
+    successfully check their keys.
     """
 
     a1 = torch.randn([2])
@@ -78,8 +78,8 @@ def test_conjunct_check_keys_1():
 
 def test_conjunct_check_keys_2():
     """
-    Tests that `check_keys` works correctly for a conjunction of transforms: their
-    `output_keys` should be disjoint.
+    Tests that `check_keys` works correctly for a conjunction of transforms: their `output_keys`
+    should be disjoint.
     """
 
     a1 = torch.randn([2])

tests/unit/autojac/_transform/test_interactions.py

@@ -253,8 +253,8 @@ def test_equivalence_jac_grads():
 
 def test_stack_check_keys():
     """
-    Tests that the `check_keys` method works correctly for a stack of transforms: all of
-    them should have the same `required_keys`.
+    Tests that the `check_keys` method works correctly for a stack of transforms: all of them should
+    have the same `required_keys`.
     """
 
     a = torch.tensor(1.0, requires_grad=True)

tests/unit/autojac/_transform/test_select.py

@@ -58,8 +58,8 @@ def test_conjunction_of_selects_is_select():
 
 def test_check_keys():
     """
-    Tests that the `check_keys` method works correctly: the set of keys to select should
-    be a subset of the set of required_keys.
+    Tests that the `check_keys` method works correctly: the set of keys to select should be a subset
+    of the set of required_keys.
     """
 
     key1 = torch.tensor([1.0])