[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

thunder/core/rematerialization.py

@@ -687,7 +687,7 @@ def add_to_swapmap(p):
     )
     new_fw_trace.bound_symbols.append(replace(fw_trace.bound_symbols[-1], args=fw_trace.bound_symbols[-1].args))
 
-    # outputs required for backward may have different names between forward and backward. 
+    # outputs required for backward may have different names between forward and backward.
     # Rematerialisation may remove some outs from the forward.
     old_saved_for_backward_fw = fw_trace.bound_symbols[-1].args[1][0]
     old_saved_for_backward_bw = []
@@ -696,8 +696,12 @@ def add_to_swapmap(p):
             old_saved_for_backward_bw = bsym.flat_outs
             break
     assert len(old_saved_for_backward_fw) == len(old_saved_for_backward_bw)
-    new_required_for_bakward_fw_to_bw_map = {x.name: y for x, y in zip(old_saved_for_backward_bw, old_saved_for_backward_fw) if x is not None}
-    new_required_for_backward = tuple([new_required_for_bakward_fw_to_bw_map[a.name] for a in new_required_for_backward])
+    new_required_for_bakward_fw_to_bw_map = {
+        x.name: y for x, y in zip(old_saved_for_backward_bw, old_saved_for_backward_fw) if x is not None
+    }
+    new_required_for_backward = tuple(
+        [new_required_for_bakward_fw_to_bw_map[a.name] for a in new_required_for_backward]
+    )
     _update_forward_with_new_saved_for_backward(new_fw_trace, new_required_for_backward)
 
     # prims.python_return was updated and now DCE can remove the unused

thunder/tests/test_transforms.py

@@ -670,9 +670,7 @@ def transform_module(self, model: thunder.ThunderModule):
                 }
                 model._overrides_buffers[n] = qb
 
-        def transform_traces_pre_prologue(
-            self, prologue_trace, computation_trace, epilogue_trace, **kwargs
-        ):
+        def transform_traces_pre_prologue(self, prologue_trace, computation_trace, epilogue_trace, **kwargs):
             tm = self.thunder_module
             from thunder.core.trace import tracectx
 
@@ -681,19 +679,15 @@ def transform_traces_pre_prologue(
             prologue_proxy_map = {
                 get_param_bsym.output.name: dict(
                     shape=self.cast_states[model_weight_name]["qb.shape"],
-                    dtype=thunder.dtypes.to_dtype(
-                        self.cast_states[model_weight_name]["qb.dtype"]
-                    ),
+                    dtype=thunder.dtypes.to_dtype(self.cast_states[model_weight_name]["qb.dtype"]),
                 )
                 for model_weight_name, (check_bsym, get_param_bsym) in checks.items()
                 if model_weight_name in self.cast_states
             }
 
             # here we switch the prologue_trace to a copy with new metadata
-            prologue_trace = (
-                thunder.transforms.quantization.trace_with_replaced_proxy_metadata(
-                    prologue_trace, prologue_proxy_map
-                )
+            prologue_trace = thunder.transforms.quantization.trace_with_replaced_proxy_metadata(
+                prologue_trace, prologue_proxy_map
             )
 
             checks = thunder.transforms.utils.get_checks(prologue_trace)
@@ -714,79 +708,58 @@ def transform_traces_pre_prologue(
                     shape=psym.shape,
                     dtype=psym.dtype,
                 )
-                for psym, csym in zip(
-                    prologue_trace.bound_symbols[-1].args[0][0], computation_trace.args
-                )
+                for psym, csym in zip(prologue_trace.bound_symbols[-1].args[0][0], computation_trace.args)
                 if psym.shape != csym.shape or psym.dtype != csym.dtype
             }
 
-            new_computation_trace = (
-                thunder.transforms.quantization.trace_with_replaced_proxy_metadata(
-                    computation_trace, computation_proxy_map
-                )
+            new_computation_trace = thunder.transforms.quantization.trace_with_replaced_proxy_metadata(
+                computation_trace, computation_proxy_map
             )
 
-            producers, consumers = thunder.core.utils.producers_and_consumers(
-                new_computation_trace
-            )
+            producers, consumers = thunder.core.utils.producers_and_consumers(new_computation_trace)
 
             bound_symbols = new_computation_trace.bound_symbols
             new_computation_trace.bound_symbols = []
 
-            new_computation_trace._siginfo.args = [
-                (a.name, None) for a in new_computation_trace.args
-            ]
+            new_computation_trace._siginfo.args = [(a.name, None) for a in new_computation_trace.args]
 
             computation_proxy_map = {}
             new_bound_symbols = []
             for bsym in bound_symbols:
-                if (
-                    bsym.sym == thunder.torch.to
-                    and producers[bsym.args[0]].sym == thunder.core.prims.unpack_trivial
-                ):
+                if bsym.sym == thunder.torch.to and producers[bsym.args[0]].sym == thunder.core.prims.unpack_trivial:
                     inp = bsym.args[0]
                     args = (inp, inp.dtype, *bsym.args[2:])
-                    computation_proxy_map[bsym.output.name] = dict(
-                        shape=inp.shape, dtype=inp.dtype
-                    )
+                    computation_proxy_map[bsym.output.name] = dict(shape=inp.shape, dtype=inp.dtype)
                     assert (
-                        len(bsym.subsymbols) == 1
-                        and bsym.subsymbols[0].sym
-                        == thunder.core.prims.convert_element_type
+                        len(bsym.subsymbols) == 1 and bsym.subsymbols[0].sym == thunder.core.prims.convert_element_type
                     )
                     subsymbols = [bsym.subsymbols[0].from_bsym(args=(inp, inp.dtype))]
-                    new_bound_symbols.append(
-                        bsym.from_bsym(args=args, subsymbols=subsymbols)
-                    )
+                    new_bound_symbols.append(bsym.from_bsym(args=args, subsymbols=subsymbols))
                 else:
                     new_bound_symbols.append(bsym.from_bsym())
 
             new_computation_trace.bound_symbols = new_bound_symbols
 
-            new_computation_trace = (
-                thunder.transforms.quantization.trace_with_replaced_proxy_metadata(
-                    new_computation_trace, computation_proxy_map
-                )
+            new_computation_trace = thunder.transforms.quantization.trace_with_replaced_proxy_metadata(
+                new_computation_trace, computation_proxy_map
             )
 
-            new_computation_trace.set_provenance(
-                thunder.core.trace.TraceProvenance("Dtype Convert")
-            )
+            new_computation_trace.set_provenance(thunder.core.trace.TraceProvenance("Dtype Convert"))
             return prologue_trace, new_computation_trace, epilogue_trace
 
     class cast(nn.Module):
         def __init__(
             self,
-            k_shape: Tuple[int, int, int, int],
-            v_shape: Tuple[int, int, int, int],
-            device: Optional[torch.device] = None,
-            dtype: Optional[torch.dtype] = None,
+            k_shape: tuple[int, int, int, int],
+            v_shape: tuple[int, int, int, int],
+            device: torch.device | None = None,
+            dtype: torch.dtype | None = None,
         ) -> None:
             super().__init__()
             self.register_buffer("k", torch.zeros(k_shape, device=device, dtype=dtype), persistent=False)
             self.register_buffer("v", torch.zeros(v_shape, device=device, dtype=dtype), persistent=False)
 
-        def forward(self, k: torch.Tensor, v: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        def forward(self, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
             # move the buffer to the activation dtype for when AMP is used
             self.k = self.k.to(k.dtype)
             self.v = self.v.to(v.dtype)
@@ -795,23 +768,34 @@ def forward(self, k: torch.Tensor, v: torch.Tensor) -> Tuple[torch.Tensor, torch
 
     # BUG: issue: 1637
     class ParentModule(nn.Module):
-        def __init__(self, k_shape: Tuple[int, int, int, int], v_shape: Tuple[int, int, int, int], device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        def __init__(
+            self,
+            k_shape: tuple[int, int, int, int],
+            v_shape: tuple[int, int, int, int],
+            device: torch.device | None = None,
+            dtype: torch.dtype | None = None,
+        ):
             super().__init__()
             self.cast_module = cast(k_shape, v_shape, device=device, dtype=dtype)
 
-        def forward(self, k: torch.Tensor, v: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        def forward(self, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
             return self.cast_module(k, v)
 
     with torch.device("cpu"):
         k_shape = (2, 3, 4, 5)
         v_shape = (2, 3, 4, 5)
         device = torch.device("cpu")
         dtype = torch.float32
-        model = (ParentModule(k_shape, v_shape, device=device, dtype=dtype).eval().requires_grad_(False))
+        model = ParentModule(k_shape, v_shape, device=device, dtype=dtype).eval().requires_grad_(False)
 
     k = torch.randn(2, 3, 4, 5, device=device, dtype=torch.half)
     v = torch.randn(2, 3, 4, 5, device=device, dtype=torch.half)
-    cast_jit = thunder.jit(model, transforms=[CastBuffers(),])
+    cast_jit = thunder.jit(
+        model,
+        transforms=[
+            CastBuffers(),
+        ],
+    )
     output_k, output_v = cast_jit(k, v)
 
     def check_dtypes(bsym):